import torch
from diffusers import DiffusionPipeline  # 기본 클래스 사용
from PIL import Image
from io import BytesIO
import base64

# 커스텀 모듈 임포트
from custom_pipeline import WanTransformer3DModel, AutoencoderKLWan

# 모델 ID
model_id = "grnr9730/Wan2.1-I2V-14B-720P-Diffusers"

# 커스텀 파이프라인 정의
class WanImageToVideoPipeline(DiffusionPipeline):
    def __init__(self, transformer, vae, scheduler):
        super().__init__()
        self.transformer = transformer
        self.vae = vae
        self.scheduler = scheduler
        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)

    def __call__(self, prompt, **kwargs):
        # 여기에 실제 추론 로직을 추가 (WanTransformer3DModel과 AutoencoderKLWan 사용)
        # 예시로 간단히 작성
        latents = self.vae.encode(torch.randn(1, 3, 224, 224)).latent_dist.sample()
        for _ in self.scheduler.timesteps:
            latents = self.transformer(latents)
        video_frames = self.vae.decode(latents).sample
        return type('Result', (), {'frames': [Image.fromarray((frame * 255).byte().cpu().numpy()) for frame in video_frames]})

# 모델 로드
pipe = WanImageToVideoPipeline.from_pretrained(
    model_id,
    transformer=WanTransformer3DModel.from_pretrained(model_id),
    vae=AutoencoderKLWan.from_pretrained(model_id),
    scheduler=FlowMatchEulerDiscreteScheduler.from_pretrained(model_id),
    torch_dtype=torch.bfloat16
)
pipe.enable_model_cpu_offload()

# Inference 함수
def infer(data):
    prompt = data.get("prompt", "A futuristic cityscape")
    
    # 비디오 생성
    video = pipe(prompt).frames  # 프레임 리스트 반환

    # 첫 번째 프레임을 이미지로 저장
    first_frame = video[0]

    # 이미지 Base64로 변환
    buffered = BytesIO()
    first_frame.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

    return {"image": img_str}