| """ |
| SkyReels-V2-DF-14B - UNLIMITED LENGTH Highest Quality Video Generation |
| Job Queue Pattern: POST returns job_id, poll /status/{job_id}, GET /result/{job_id} |
| 14B model on A100 80GB - supports 540P and 720P, videos up to 60s+ |
| """ |
| import os |
| import sys |
| import subprocess |
| import base64 |
| import tempfile |
| import traceback |
| import threading |
| import time |
| import uuid |
|
|
| |
| 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-14b] torch: {torch.__version__}, CUDA: {torch.cuda.is_available()}", flush=True) |
| if torch.cuda.is_available(): |
| print(f"[video-14b] GPU: {torch.cuda.get_device_name(0)}", flush=True) |
| print(f"[video-14b] VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB", flush=True) |
|
|
| app = Flask(__name__) |
| pipe = None |
| jobs = {} |
|
|
| |
| RESOLUTION = os.environ.get("MODEL_RESOLUTION", "540P") |
| if RESOLUTION == "720P": |
| MODEL_ID = "Skywork/SkyReels-V2-DF-14B-720P-Diffusers" |
| DEFAULT_HEIGHT = 720 |
| DEFAULT_WIDTH = 1280 |
| DEFAULT_BASE_FRAMES = 121 |
| else: |
| MODEL_ID = "Skywork/SkyReels-V2-DF-14B-540P-Diffusers" |
| DEFAULT_HEIGHT = 544 |
| DEFAULT_WIDTH = 960 |
| DEFAULT_BASE_FRAMES = 97 |
|
|
| print(f"[video-14b] Using model: {MODEL_ID} ({RESOLUTION})", flush=True) |
|
|
|
|
| def load_model(): |
| global pipe |
| from diffusers import ( |
| AutoModel, |
| SkyReelsV2DiffusionForcingPipeline, |
| UniPCMultistepScheduler, |
| ) |
| print(f"[video-14b] Loading {MODEL_ID}...", flush=True) |
|
|
| vae = AutoModel.from_pretrained( |
| MODEL_ID, |
| subfolder="vae", |
| torch_dtype=torch.float32, |
| ) |
|
|
| pipe = SkyReelsV2DiffusionForcingPipeline.from_pretrained( |
| MODEL_ID, |
| vae=vae, |
| torch_dtype=torch.bfloat16, |
| ) |
|
|
| pipe.scheduler = UniPCMultistepScheduler.from_config( |
| pipe.scheduler.config, flow_shift=8.0 |
| ) |
| pipe.to("cuda") |
| print(f"[video-14b] Model loaded on CUDA! ({RESOLUTION})", flush=True) |
|
|
|
|
| def generate_video_background(job_id, prompt, params): |
| """Run video generation in a background thread.""" |
| global jobs |
| try: |
| from diffusers.utils import export_to_video |
|
|
| jobs[job_id]["status"] = "generating" |
| jobs[job_id]["progress"] = "Starting generation..." |
|
|
| num_frames = int(params.get("num_frames", 257)) |
| height = int(params.get("height", DEFAULT_HEIGHT)) |
| width = int(params.get("width", DEFAULT_WIDTH)) |
| steps = int(params.get("num_inference_steps", 10)) |
| fps = int(params.get("fps", 24)) |
| base_num_frames = int(params.get("base_num_frames", DEFAULT_BASE_FRAMES)) |
| ar_step = int(params.get("ar_step", 5)) |
| 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", 6.0)) |
|
|
| duration_sec = num_frames / fps |
| jobs[job_id]["progress"] = f"Generating {num_frames} frames ({duration_sec:.1f}s) at {RESOLUTION}..." |
| print(f"[video-14b] Job {job_id}: {num_frames} frames ({duration_sec:.1f}s), {prompt[:80]}", flush=True) |
|
|
| 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, |
| guidance_scale=guidance_scale, |
| ) |
| if ar_step > 0: |
| call_kwargs["causal_block_size"] = causal_block_size |
|
|
| output = pipe(**call_kwargs).frames[0] |
|
|
| jobs[job_id]["progress"] = "Encoding video..." |
|
|
| 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 |
|
|
| jobs[job_id]["status"] = "completed" |
| jobs[job_id]["video_b64"] = video_b64 |
| jobs[job_id]["frames"] = len(output) |
| jobs[job_id]["duration"] = duration |
| jobs[job_id]["progress"] = f"Done: {len(output)} frames, {duration:.1f}s" |
| print(f"[video-14b] Job {job_id} completed: {len(output)} frames, {duration:.1f}s", flush=True) |
|
|
| except Exception as e: |
| tb = traceback.format_exc() |
| jobs[job_id]["status"] = "failed" |
| jobs[job_id]["error"] = str(e) |
| jobs[job_id]["traceback"] = tb |
| print(f"[video-14b] Job {job_id} FAILED: {e}", flush=True) |
| print(f"[video-14b] Traceback:\n{tb}", flush=True) |
|
|
|
|
| @app.route("/health", methods=["GET"]) |
| def health(): |
| if pipe is not None: |
| return jsonify({"status": "healthy", "model": "skyreels-v2-df-14b", "resolution": RESOLUTION}), 200 |
| return jsonify({"status": "loading"}), 503 |
|
|
|
|
| @app.route("/", methods=["POST"]) |
| @app.route("/api/generate-video", methods=["POST"]) |
| def submit_video_job(): |
| """Submit a video generation job. Returns job_id immediately.""" |
| try: |
| data = request.get_json(force=True) |
| prompt = data.get("inputs", "") or data.get("prompt", "") |
| params = data.get("parameters", {}) |
|
|
| if not prompt: |
| return jsonify({"error": "Prompt is required", "success": False}), 400 |
|
|
| job_id = str(uuid.uuid4())[:8] |
| jobs[job_id] = { |
| "status": "queued", |
| "progress": "Queued for generation", |
| "video_b64": None, |
| "frames": None, |
| "duration": None, |
| "error": None, |
| "prompt": prompt, |
| "params": params, |
| "created_at": time.time(), |
| } |
|
|
| thread = threading.Thread( |
| target=generate_video_background, |
| args=(job_id, prompt, params), |
| daemon=True |
| ) |
| thread.start() |
|
|
| print(f"[video-14b] Job {job_id} submitted: {prompt[:80]}", flush=True) |
|
|
| return jsonify({ |
| "success": True, |
| "job_id": job_id, |
| "status": "queued", |
| "message": f"Video generation started. Poll /status/{job_id} for progress.", |
| "status_url": f"/status/{job_id}", |
| "result_url": f"/result/{job_id}", |
| "model": "skyreels-v2-df-14b", |
| "resolution": RESOLUTION, |
| }) |
|
|
| except Exception as e: |
| return jsonify({"error": str(e), "success": False}), 500 |
|
|
|
|
| @app.route("/status/<job_id>", methods=["GET"]) |
| def job_status(job_id): |
| if job_id not in jobs: |
| return jsonify({"error": "Job not found", "success": False}), 404 |
|
|
| job = jobs[job_id] |
| return jsonify({ |
| "success": True, |
| "job_id": job_id, |
| "status": job["status"], |
| "progress": job["progress"], |
| "frames": job["frames"], |
| "duration": job["duration"], |
| "error": job["error"], |
| }) |
|
|
|
|
| @app.route("/result/<job_id>", methods=["GET"]) |
| def job_result(job_id): |
| if job_id not in jobs: |
| return jsonify({"error": "Job not found", "success": False}), 404 |
|
|
| job = jobs[job_id] |
| if job["status"] != "completed": |
| return jsonify({ |
| "success": False, |
| "status": job["status"], |
| "progress": job["progress"], |
| "message": "Video not ready yet. Keep polling /status/" + job_id, |
| }), 202 |
|
|
| return jsonify({ |
| "success": True, |
| "job_id": job_id, |
| "video": job["video_b64"], |
| "video_base64": job["video_b64"], |
| "format": "mp4", |
| "frames": job["frames"], |
| "duration": job["duration"], |
| "duration_seconds": job["duration"], |
| "model": "skyreels-v2-df-14b", |
| "resolution": RESOLUTION, |
| }) |
|
|
|
|
| @app.route("/debug", methods=["GET"]) |
| def debug(): |
| return jsonify({ |
| "pipe_loaded": pipe is not None, |
| "model": "skyreels-v2-df-14b", |
| "resolution": RESOLUTION, |
| "model_id": MODEL_ID, |
| "active_jobs": {k: {"status": v["status"], "progress": v["progress"]} for k, v in jobs.items()}, |
| "total_jobs": len(jobs), |
| }) |
|
|
|
|
| if __name__ == "__main__": |
| load_model() |
| app.run(host="0.0.0.0", port=8000, threaded=True) |
|
|
