Upload server.py with huggingface_hub

server.py

@@ -0,0 +1,86 @@
+import os
+import sys
+import subprocess
+import base64
+import io
+import json
+
+# Install diffusers from source (supports WanPipeline)
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "git+https://github.com/huggingface/diffusers.git"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "accelerate", "sentencepiece", "protobuf", "imageio[ffmpeg]"])
+
+import torch
+from flask import Flask, request, jsonify
+from diffusers import AutoencoderKLWan, WanPipeline
+from diffusers.utils import export_to_video
+
+app = Flask(__name__)
+pipe = None
+
+def load_model():
+    global pipe
+    print("[video] Loading Wan2.2-TI2V-5B pipeline...", flush=True)
+    vae = AutoencoderKLWan.from_pretrained(
+        "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+        subfolder="vae",
+        torch_dtype=torch.float32,
+    )
+    pipe = WanPipeline.from_pretrained(
+        "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+        vae=vae,
+        torch_dtype=torch.bfloat16,
+    )
+    pipe.to("cuda")
+    print("[video] Model loaded successfully!", flush=True)
+
+@app.route("/health", methods=["GET"])
+def health():
+    if pipe is not None:
+        return jsonify({"status": "healthy"}), 200
+    return jsonify({"status": "loading"}), 503
+
+@app.route("/", methods=["POST"])
+def generate():
+    try:
+        data = request.get_json()
+        prompt = data.get("inputs", "")
+        params = data.get("parameters", {})
+
+        num_frames = int(params.get("num_frames", 49))
+        height = int(params.get("height", 480))
+        width = int(params.get("width", 832))
+        steps = int(params.get("num_inference_steps", 30))
+        fps = int(params.get("fps", 24))
+        guidance = float(params.get("guidance_scale", 5.0))
+        negative = params.get("negative_prompt", "low quality, blurry, distorted")
+
+        print(f"[video] Generating {num_frames} frames: {prompt[:100]}...", flush=True)
+        frames = pipe(
+            prompt=prompt,
+            negative_prompt=negative,
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+        ).frames[0]
+
+        buf = io.BytesIO()
+        export_to_video(frames, buf, fps=fps)
+        video_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        duration = len(frames) / fps
+        print(f"[video] Done, {len(frames)} frames, {duration:.1f}s", flush=True)
+
+        return jsonify({
+            "video": video_b64,
+            "format": "mp4",
+            "frames": len(frames),
+            "duration": duration,
+        })
+    except Exception as e:
+        print(f"[video] Error: {e}", flush=True)
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == "__main__":
+    load_model()
+    app.run(host="0.0.0.0", port=8000)