File size: 4,088 Bytes
dcb6cc1 afe2d60 eb97492 dcb6cc1 afe2d60 eb97492 dcb6cc1 eb97492 04d3033 dcb6cc1 eb97492 dcb6cc1 eb97492 dcb6cc1 afe2d60 dcb6cc1 afe2d60 dcb6cc1 eb97492 06447aa dcb6cc1 eb97492 dcb6cc1 eb97492 dcb6cc1 eb97492 dcb6cc1 afe2d60 dcb6cc1 afe2d60 dcb6cc1 06447aa dcb6cc1 eb97492 06447aa dcb6cc1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import os
import sys
import subprocess
import base64
import tempfile
import traceback
# Install latest diffusers + dependencies
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade",
"diffusers", "flask", "accelerate", "sentencepiece", "protobuf", "imageio[ffmpeg]", "transformers", "huggingface_hub"])
import torch
from flask import Flask, request, jsonify
print(f"[video] torch version: {torch.__version__}", flush=True)
print(f"[video] CUDA available: {torch.cuda.is_available()}", flush=True)
if torch.cuda.is_available():
print(f"[video] GPU: {torch.cuda.get_device_name(0)}", flush=True)
print(f"[video] VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB", flush=True)
app = Flask(__name__)
pipe = None
def load_model():
global pipe
from diffusers import AutoencoderKLWan, WanPipeline
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.bfloat16,
)
pipe = WanPipeline.from_pretrained(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
vae=vae,
torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
print("[video] Model loaded with CPU offload!", 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:
from diffusers.utils import export_to_video
data = request.get_json(force=True)
# Support both HF format (inputs) and backend format (prompt)
prompt = data.get("inputs", "") or data.get("prompt", "")
params = data.get("parameters", {})
num_frames = int(params.get("num_frames", 25))
height = int(params.get("height", 480))
width = int(params.get("width", 832))
steps = int(params.get("num_inference_steps", 15))
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)
print(f"[video] Params: {width}x{height}, steps={steps}, fps={fps}", 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]
# 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(video_bytes)} bytes", flush=True)
return jsonify({
"video": video_b64,
"format": "mp4",
"frames": len(frames),
"duration": duration,
"success": True,
"video_base64": video_b64,
"duration_seconds": duration,
"model": "wan2.2-ti2v-5b",
})
except Exception as e:
tb = traceback.format_exc()
print(f"[video] ERROR: {e}\n{tb}", flush=True)
return jsonify({"error": str(e), "traceback": tb, "success": False}), 500
# Also add /api/generate-video route for backend compatibility
@app.route("/api/generate-video", methods=["POST"])
def api_generate_video():
"""Handle the backend's /api/generate-video format by forwarding to generate()"""
return generate()
if __name__ == "__main__":
load_model()
app.run(host="0.0.0.0", port=8000)
|