File size: 5,837 Bytes
f2263c7 b8a4c12 f2263c7 32c534e f2263c7 4e63a35 f2263c7 5540851 f2263c7 4e63a35 f2263c7 4e63a35 f2263c7 5540851 f2263c7 bacb366 f2263c7 5540851 f2263c7 | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | import os
import sys
import subprocess
import base64
import tempfile
import traceback
# Install latest diffusers (supports SkyReelsV2DiffusionForcingPipeline) + dependencies
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade",
"diffusers", "flask", "accelerate", "sentencepiece", "protobuf", "imageio[ffmpeg]",
"transformers", "huggingface_hub", "ftfy", "einops"])
import torch
from flask import Flask, request, jsonify
print(f"[video-long] torch version: {torch.__version__}", flush=True)
print(f"[video-long] CUDA available: {torch.cuda.is_available()}", flush=True)
if torch.cuda.is_available():
print(f"[video-long] GPU: {torch.cuda.get_device_name(0)}", flush=True)
print(f"[video-long] 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 (
AutoModel,
SkyReelsV2DiffusionForcingPipeline,
UniPCMultistepScheduler,
)
print("[video-long] Loading SkyReels-V2-DF-1.3B-540P...", flush=True)
model_id = "Skywork/SkyReels-V2-DF-1.3B-540P-Diffusers"
vae = AutoModel.from_pretrained(
model_id,
subfolder="vae",
torch_dtype=torch.float32,
)
pipe = SkyReelsV2DiffusionForcingPipeline.from_pretrained(
model_id,
vae=vae,
torch_dtype=torch.bfloat16,
)
# Flow shift 8.0 for T2V
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=8.0
)
pipe.to("cuda")
print("[video-long] Model loaded on CUDA!", 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("/debug", methods=["GET"])
def debug():
if pipe is None:
return jsonify({"status": "loading"}), 503
import inspect
sig = inspect.signature(pipe.__call__)
params = {}
for name, param in sig.parameters.items():
params[name] = str(param.default) if param.default != inspect.Parameter.empty else "required"
return jsonify({
"pipeline_class": type(pipe).__name__,
"parameters": params,
"device": str(pipe.device) if hasattr(pipe, 'device') else 'unknown',
})
@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", {})
# Diffusion Forcing params - defaults for ~10s video
num_frames = int(params.get("num_frames", 257)) # 257=10s, 737=30s, 1457=60s
height = int(params.get("height", 544)) # 544 for 540P
width = int(params.get("width", 960)) # 960 for 540P
steps = int(params.get("num_inference_steps", 30))
fps = int(params.get("fps", 24))
base_num_frames = int(params.get("base_num_frames", 97))
ar_step = int(params.get("ar_step", 5)) # 5 for async (better quality, faster for long)
causal_block_size = int(params.get("causal_block_size", 5))
overlap_history = int(params.get("overlap_history", 17))
addnoise_condition = int(params.get("addnoise_condition", 20))
guidance_scale = float(params.get("guidance_scale", 5.0))
print(f"[video-long] Generating {num_frames} frames ({num_frames/fps:.1f}s): {prompt[:100]}", flush=True)
print(f"[video-long] Params: {width}x{height}, steps={steps}, ar_step={ar_step}", flush=True)
# Build kwargs - only pass causal_block_size when ar_step > 0
call_kwargs = dict(
prompt=prompt,
num_inference_steps=steps,
height=height,
width=width,
num_frames=num_frames,
base_num_frames=base_num_frames,
ar_step=ar_step,
overlap_history=overlap_history,
addnoise_condition=addnoise_condition,
)
if ar_step > 0:
call_kwargs["causal_block_size"] = causal_block_size
print(f"[video-long] Pipeline call kwargs: {call_kwargs}", flush=True)
output = pipe(**call_kwargs).frames[0]
# Write to temp file
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
tmp_path = tmp.name
export_to_video(output, 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(output) / fps
print(f"[video-long] Done: {len(output)} frames, {duration:.1f}s, {len(video_bytes)} bytes", flush=True)
return jsonify({
"video": video_b64,
"format": "mp4",
"frames": len(output),
"duration": duration,
"success": True,
"video_base64": video_b64,
"duration_seconds": duration,
"model": "skyreels-v2-df-1.3b",
})
except Exception as e:
tb = traceback.format_exc()
print(f"[video-long] ERROR: {e}", flush=True)
print(f"[video-long] TRACEBACK:\n{tb}", flush=True)
return jsonify({"error": str(e), "traceback": tb, "success": False}), 200
# 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)
|