Upload server.py with huggingface_hub

server.py

@@ -2,10 +2,10 @@ import os
 import sys
 import subprocess
 import base64
-import io
+import tempfile
 import traceback
 
-# Install latest diffusers (supports WanPipeline) + dependencies
+# Install latest diffusers + dependencies
 subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade",
     "diffusers", "flask", "accelerate", "sentencepiece", "protobuf", "imageio[ffmpeg]", "transformers", "huggingface_hub"])
 
@@ -26,19 +26,18 @@ def load_model():
     from diffusers import AutoencoderKLWan, WanPipeline
     print("[video] Loading Wan2.2-TI2V-5B pipeline...", flush=True)
 
-    # VAE in float32 for quality, model in bfloat16 for speed
     vae = AutoencoderKLWan.from_pretrained(
         "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
         subfolder="vae",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.bfloat16,
     )
     pipe = WanPipeline.from_pretrained(
         "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
         vae=vae,
         torch_dtype=torch.bfloat16,
     )
-    pipe.to("cuda")
-    print("[video] Model loaded on CUDA!", flush=True)
+    pipe.enable_model_cpu_offload()
+    print("[video] Model loaded with CPU offload!", flush=True)
 
 @app.route("/health", methods=["GET"])
 def health():
@@ -76,11 +75,19 @@ def generate():
             guidance_scale=guidance,
         ).frames[0]
 
-        buf = io.BytesIO()
-        export_to_video(frames, buf, fps=fps)
-        video_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        # Write to temp file (export_to_video needs a file path, not BytesIO)
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+            tmp_path = tmp.name
+
+        export_to_video(frames, tmp_path, fps=fps)
+
+        with open(tmp_path, "rb") as f:
+            video_bytes = f.read()
+        os.unlink(tmp_path)
+
+        video_b64 = base64.b64encode(video_bytes).decode("utf-8")
         duration = len(frames) / fps
-        print(f"[video] Done: {len(frames)} frames, {duration:.1f}s, {len(buf.getvalue())} bytes", flush=True)
+        print(f"[video] Done: {len(frames)} frames, {duration:.1f}s, {len(video_bytes)} bytes", flush=True)
 
         return jsonify({
             "video": video_b64,
@@ -90,7 +97,8 @@ def generate():
         })
     except Exception as e:
         tb = traceback.format_exc()
-        print(f"[video] ERROR: {e}\n{tb}", flush=True)
+        print(f"[video] ERROR: {e}", flush=True)
+        print(f"[video] TRACEBACK:\n{tb}", flush=True)
         return jsonify({"error": str(e), "traceback": tb}), 500
 
 if __name__ == "__main__":