Update handler.py

handler.py

@@ -1,26 +1,29 @@
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import WanPipeline
 from PIL import Image
 from io import BytesIO
 import base64
 
-# 모델 로드
+# 모델 ID
 model_id = "grnr9730/Wan2.1-I2V-14B-720P-Diffusers"
-#pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
-pipe = DiffusionPipeline.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float16, 
-    custom_pipeline="Wan2.1-I2V-14B-720P-Diffusers"
-).to("cuda")
+
+# 모델 로드
+pipe = WanPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload()
 
 # Inference 함수
 def infer(data):
     prompt = data.get("prompt", "A futuristic cityscape")
-    image = pipe(prompt).images[0]
     
+    # 비디오 생성
+    video = pipe(prompt).frames  # 프레임 리스트 반환
+
+    # 첫 번째 프레임을 이미지로 저장
+    first_frame = video[0]
+
     # 이미지 Base64로 변환
     buffered = BytesIO()
-    image.save(buffered, format="PNG")
+    first_frame.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    
-    return {"image": img_str}
+
+    return {"image": img_str}