Upload folder using huggingface_hub

handler.py

@@ -0,0 +1,58 @@
+import base64
+import io
+import torch
+from diffusers import AutoencoderKLWan, WanPipeline
+from diffusers.utils import export_to_video
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.bfloat16
+
+        # Load VAE and pipeline from the diffusers-format repo
+        vae = AutoencoderKLWan.from_pretrained(
+            "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+            subfolder="vae",
+            torch_dtype=torch.float32,
+        )
+        self.pipe = WanPipeline.from_pretrained(
+            "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+            vae=vae,
+            torch_dtype=self.dtype,
+        )
+        self.pipe.to(self.device)
+
+    def __call__(self, data):
+        prompt = data.pop("inputs", "")
+        params = data.pop("parameters", {})
+
+        num_frames = params.get("num_frames", 49)
+        height = params.get("height", 480)
+        width = params.get("width", 832)
+        num_inference_steps = params.get("num_inference_steps", 30)
+        fps = params.get("fps", 24)
+        guidance_scale = params.get("guidance_scale", 5.0)
+        negative_prompt = params.get("negative_prompt", "low quality, blurry, distorted")
+
+        frames = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+        ).frames[0]
+
+        # Export to video in memory
+        buf = io.BytesIO()
+        export_to_video(frames, buf, fps=fps)
+        video_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+
+        duration = len(frames) / fps
+        return {
+            "video": video_b64,
+            "format": "mp4",
+            "frames": len(frames),
+            "duration": duration,
+        }

requirements.txt

@@ -0,0 +1,7 @@
+git+https://github.com/huggingface/diffusers.git
+transformers>=4.51.0
+accelerate>=0.34.0
+torch>=2.1.0
+sentencepiece
+protobuf
+imageio[ffmpeg]