import os import gc import time import random import torch import imageio import gradio as gr from diffusers.utils import load_image from skyreels_v2_infer import DiffusionForcingPipeline from skyreels_v2_infer.modules import download_model from skyreels_v2_infer.pipelines import PromptEnhancer, resizecrop def generate_diffusion_forced_video( prompt, model_id, resolution, num_frames, image=None, ar_step=0, causal_attention=False, causal_block_size=1, base_num_frames=97, overlap_history=None, addnoise_condition=0, guidance_scale=6.0, shift=8.0, inference_steps=30, use_usp=False, offload=True, fps=24, seed=None, prompt_enhancer=False, teacache=False, teacache_thresh=0.2, use_ret_steps=False ): model_id = download_model(model_id) if resolution == "540P": height, width = 544, 960 elif resolution == "720P": height, width = 720, 1280 else: raise ValueError(f"Invalid resolution: {resolution}") if seed is None: random.seed(time.time()) seed = int(random.randrange(4294967294)) if num_frames > base_num_frames and overlap_history is None: raise ValueError("Specify `overlap_history` for long video generation. Try 17 or 37.") if addnoise_condition > 60: print("Warning: Large `addnoise_condition` may reduce consistency. Recommended: 20.") if image is not None: image = load_image(image).convert("RGB") image_width, image_height = image.size if image_height > image_width: height, width = width, height image = resizecrop(image, height, width) negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" prompt_input = prompt if prompt_enhancer and image is None: enhancer = PromptEnhancer() prompt_input = enhancer(prompt_input) del enhancer gc.collect() torch.cuda.empty_cache() pipe = DiffusionForcingPipeline( model_id, dit_path=model_id, device=torch.device("cuda"), weight_dtype=torch.bfloat16, use_usp=use_usp, offload=offload, ) if causal_attention: pipe.transformer.set_ar_attention(causal_block_size) if teacache: if ar_step > 0: num_steps = ( inference_steps + (((base_num_frames - 1) // 4 + 1) // causal_block_size - 1) * ar_step ) else: num_steps = inference_steps pipe.transformer.initialize_teacache( enable_teacache=True, num_steps=num_steps, teacache_thresh=teacache_thresh, use_ret_steps=use_ret_steps, ckpt_dir=model_id, ) with torch.amp.autocast("cuda", dtype=pipe.transformer.dtype), torch.no_grad(): video_frames = pipe( prompt=prompt_input, negative_prompt=negative_prompt, image=image, height=height, width=width, num_frames=num_frames, num_inference_steps=inference_steps, shift=shift, guidance_scale=guidance_scale, generator=torch.Generator(device="cuda").manual_seed(seed), overlap_history=overlap_history, addnoise_condition=addnoise_condition, base_num_frames=base_num_frames, ar_step=ar_step, causal_block_size=causal_block_size, fps=fps, )[0] os.makedirs("gradio_df_videos", exist_ok=True) timestamp = time.strftime("%Y%m%d_%H%M%S") output_path = f"gradio_df_videos/{prompt[:50].replace('/', '')}_{seed}_{timestamp}.mp4" imageio.mimwrite(output_path, video_frames, fps=fps, quality=8, output_params=["-loglevel", "error"]) return output_path # Gradio UI resolution_options = ["540P", "720P"] model_options = ["Skywork/SkyReels-V2-DF-1.3B-540P"] # Update if there are more with gr.Blocks() as demo: with gr.Column(): gr.Markdown("# SkyReels V2") with gr.Row(): with gr.Column(): prompt = gr.Textbox(label="Prompt") model_id = gr.Dropdown(choices=model_options, value=model_options[0], label="Model ID") resolution = gr.Radio(choices=resolution_options, value="540P", label="Resolution", interactive=False) num_frames = gr.Slider(minimum=16, maximum=200, value=97, step=1, label="Number of Frames") image = gr.Image(type="filepath", label="Input Image (optional)") with gr.Accordion("Advanced Settings", open=False): ar_step = gr.Number(label="AR Step", value=0) causal_attention = gr.Checkbox(label="Causal Attention") causal_block_size = gr.Number(label="Causal Block Size", value=1) base_num_frames = gr.Number(label="Base Num Frames", value=97) overlap_history = gr.Number(label="Overlap History (set for long videos)", value=None) addnoise_condition = gr.Number(label="AddNoise Condition", value=0) guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, value=6.0, step=0.1, label="Guidance Scale") shift = gr.Slider(minimum=0.0, maximum=20.0, value=8.0, step=0.1, label="Shift") inference_steps = gr.Slider(minimum=1, maximum=100, value=30, step=1, label="Inference Steps") use_usp = gr.Checkbox(label="Use USP") offload = gr.Checkbox(label="Offload", value=True, interactive=False) fps = gr.Slider(minimum=1, maximum=60, value=24, step=1, label="FPS") seed = gr.Number(label="Seed (optional)", precision=0) prompt_enhancer = gr.Checkbox(label="Prompt Enhancer") use_teacache = gr.Checkbox(label="Use TeaCache") teacache_thresh = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.01, label="TeaCache Threshold") use_ret_steps = gr.Checkbox(label="Use Retention Steps") submit_btn = gr.Button("Generate") with gr.Column(): output_video = gr.Video(label="Generated Video") submit_btn.click( fn = generate_diffusion_forced_video, inputs = [ prompt, model_id, resolution, num_frames, image, ar_step, causal_attention, causal_block_size, base_num_frames, overlap_history, addnoise_condition, guidance_scale, shift, inference_steps, use_usp, offload, fps, seed, prompt_enhancer, use_teacache, teacache_thresh, use_ret_steps ], outputs = [ output_video ] ) demo.launch(show_error=True, show_api=False, share=False)