Spaces:
Running
on
Zero
Running
on
Zero
import torch | |
from diffusers import AutoencoderKLWan, WanVACEPipeline UniPCMultistepScheduler | |
from diffusers.utils import export_to_video | |
from transformers import CLIPVisionModel | |
import gradio as gr | |
import tempfile | |
import spaces | |
from huggingface_hub import hf_hub_download | |
import numpy as np | |
from PIL import Image | |
import random | |
MODEL_ID = "Wan-AI/Wan2.1-VACE-14B-diffusers" | |
image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32) | |
vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32) | |
pipe = WanVACEPipeline.from_pretrained( | |
MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16 | |
) | |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0) | |
pipe.to("cuda") | |
pipe.load_lora_weights( | |
"vrgamedevgirl84/Wan14BT2VFusioniX", | |
weight_name="FusionX_LoRa/Phantom_Wan_14B_FusionX_LoRA.safetensors", | |
adapter_name="phantom" | |
) | |
pipe.load_lora_weights( | |
"vrgamedevgirl84/Wan14BT2VFusioniX", | |
weight_name="OtherLoRa's/DetailEnhancerV1.safetensors", adapter_name="detailer" | |
) | |
pipe.set_adapters(["phantom","detailer"], adapter_weights=[1, .9]) | |
pipe.fuse_lora() | |
MOD_VALUE = 32 | |
DEFAULT_H_SLIDER_VALUE = 512 | |
DEFAULT_W_SLIDER_VALUE = 896 | |
NEW_FORMULA_MAX_AREA = 480.0 * 832.0 | |
SLIDER_MIN_H, SLIDER_MAX_H = 128, 896 | |
SLIDER_MIN_W, SLIDER_MAX_W = 128, 896 | |
MAX_SEED = np.iinfo(np.int32).max | |
FIXED_FPS = 24 | |
MIN_FRAMES_MODEL = 8 | |
MAX_FRAMES_MODEL = 81 | |
# Default prompts for different modes | |
MODE_PROMPTS = { | |
"Ref2V": "", | |
"FLF2V": "", | |
"Random2V": "" | |
} | |
default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards, watermark, text, signature" | |
def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area, | |
min_slider_h, max_slider_h, | |
min_slider_w, max_slider_w, | |
default_h, default_w): | |
orig_w, orig_h = pil_image.size | |
if orig_w <= 0 or orig_h <= 0: | |
return default_h, default_w | |
aspect_ratio = orig_h / orig_w | |
calc_h = round(np.sqrt(calculation_max_area * aspect_ratio)) | |
calc_w = round(np.sqrt(calculation_max_area / aspect_ratio)) | |
calc_h = max(mod_val, (calc_h // mod_val) * mod_val) | |
calc_w = max(mod_val, (calc_w // mod_val) * mod_val) | |
new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val)) | |
new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val)) | |
return new_h, new_w | |
def handle_gallery_upload_for_dims_wan(gallery_images, current_h_val, current_w_val): | |
if gallery_images is None or len(gallery_images) == 0: | |
return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE) | |
try: | |
# Use the first image to calculate dimensions | |
first_image = gallery_images[0] | |
new_h, new_w = _calculate_new_dimensions_wan( | |
first_image, MOD_VALUE, NEW_FORMULA_MAX_AREA, | |
SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W, | |
DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE | |
) | |
return gr.update(value=new_h), gr.update(value=new_w) | |
except Exception as e: | |
gr.Warning("Error attempting to calculate new dimensions") | |
return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE) | |
def update_prompt_from_mode(mode): | |
"""Update the prompt based on the selected mode""" | |
return MODE_PROMPTS.get(mode, "") | |
def process_images_for_mode(images, mode): | |
"""Process images based on the selected mode""" | |
if not images or len(images) == 0: | |
return None | |
if mode == "Ref2V": | |
# Use the first image as reference | |
return images[0] | |
elif mode == "FLF2V": | |
# First and Last Frame: blend or interpolate between first and last image | |
if len(images) >= 2: | |
return None | |
else: | |
return images[0] | |
elif mode == "Random2V": | |
# Randomly select one image from the gallery | |
return images[0] | |
return images[0] | |
def get_duration(gallery_images, mode, prompt, height, width, | |
negative_prompt, duration_seconds, | |
guidance_scale, steps, | |
seed, randomize_seed, | |
progress): | |
if steps > 4 and duration_seconds > 2: | |
return 90 | |
elif steps > 4 or duration_seconds > 2: | |
return 75 | |
else: | |
return 60 | |
def generate_video(gallery_images, mode, prompt, height, width, | |
negative_prompt=default_negative_prompt, duration_seconds = 2, | |
guidance_scale = 1, steps = 4, | |
seed = 42, randomize_seed = False, | |
progress=gr.Progress(track_tqdm=True)): | |
""" | |
Generate a video from gallery images using the selected mode. | |
Args: | |
gallery_images (list): List of PIL images from the gallery | |
mode (str): Processing mode - "Ref2V", "FLF2V", or "Random2V" | |
prompt (str): Text prompt describing the desired animation | |
height (int): Target height for the output video | |
width (int): Target width for the output video | |
negative_prompt (str): Negative prompt to avoid unwanted elements | |
duration_seconds (float): Duration of the generated video in seconds | |
guidance_scale (float): Controls adherence to the prompt | |
steps (int): Number of inference steps | |
seed (int): Random seed for reproducible results | |
randomize_seed (bool): Whether to use a random seed | |
progress (gr.Progress): Gradio progress tracker | |
Returns: | |
tuple: (video_path, current_seed) | |
""" | |
if gallery_images is None or len(gallery_images) == 0: | |
raise gr.Error("Please upload at least one image to the gallery.") | |
# Process images based on the selected mode | |
input_image = process_images_for_mode(gallery_images, mode) | |
if input_image is None: | |
raise gr.Error("Failed to process images for the selected mode.") | |
target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE) | |
target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE) | |
num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL) | |
current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) | |
resized_image = input_image.resize((target_w, target_h)) | |
# Mode-specific processing can be added here if needed | |
if mode == "FLF2V" and len(gallery_images) >= 2: | |
# You can add special handling for FLF2V mode here | |
# For example, use both first and last frames in some way | |
pass | |
with torch.inference_mode(): | |
output_frames_list = pipe( | |
image=resized_image, prompt=prompt, negative_prompt=negative_prompt, | |
height=target_h, width=target_w, num_frames=num_frames, | |
guidance_scale=float(guidance_scale), num_inference_steps=int(steps), | |
generator=torch.Generator(device="cuda").manual_seed(current_seed) | |
).frames[0] | |
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile: | |
video_path = tmpfile.name | |
export_to_video(output_frames_list, video_path, fps=FIXED_FPS) | |
return video_path, current_seed | |
with gr.Blocks() as demo: | |
gr.Markdown("# Fast 4 steps Wan 2.1 I2V (14B) with CausVid LoRA - Multi-Image Gallery") | |
gr.Markdown("[CausVid](https://github.com/tianweiy/CausVid) is a distilled version of Wan 2.1 to run faster in just 4-8 steps, [extracted as LoRA by Kijai](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/Wan21_CausVid_14B_T2V_lora_rank32.safetensors) and is compatible with 🧨 diffusers") | |
with gr.Row(): | |
with gr.Column(): | |
# Gallery component for multiple image upload | |
gallery_component = gr.Gallery( | |
label="Upload Images", | |
show_label=True, | |
elem_id="gallery", | |
columns=3, | |
rows=2, | |
object_fit="contain", | |
height="auto", | |
type="pil", | |
allow_preview=True | |
) | |
# Radio button for mode selection | |
mode_radio = gr.Radio( | |
choices=["Ref2V", "FLF2V", "Random2V"], | |
value="Ref2V", | |
label="Processing Mode", | |
info="Ref2V: Reference to Video | FLF2V: First-Last Frame to Video | Random2V: Random Image to Video" | |
) | |
prompt_input = gr.Textbox(label="Prompt", value=MODE_PROMPTS["Ref2V"]) | |
duration_seconds_input = gr.Slider(minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), maximum=round(MAX_FRAMES_MODEL/FIXED_FPS,1), step=0.1, value=2, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.") | |
with gr.Accordion("Advanced Settings", open=False): | |
negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3) | |
seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True) | |
randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True) | |
with gr.Row(): | |
height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})") | |
width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})") | |
steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps") | |
guidance_scale_input = gr.Slider(minimum=0.0, maximum=5.0, step=0.5, value=1.0, label="Guidance Scale", visible=False) | |
generate_button = gr.Button("Generate Video", variant="primary") | |
with gr.Column(): | |
video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False) | |
with gr.Accordion("Mode Information", open=True): | |
gr.Markdown(""" | |
**Processing Modes:** | |
- **Ref2V**: Uses the first image as reference for video generation | |
- **FLF2V**: Blends first and last images for interpolation (requires at least 2 images) | |
- **Random2V**: Randomly selects one image from the gallery for generation | |
""") | |
# Update prompt when mode changes | |
mode_radio.change( | |
fn=update_prompt_from_mode, | |
inputs=[mode_radio], | |
outputs=[prompt_input] | |
) | |
# Update dimensions when gallery changes | |
gallery_component.change( | |
fn=handle_gallery_upload_for_dims_wan, | |
inputs=[gallery_component, height_input, width_input], | |
outputs=[height_input, width_input] | |
) | |
ui_inputs = [ | |
gallery_component, mode_radio, prompt_input, height_input, width_input, | |
negative_prompt_input, duration_seconds_input, | |
guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox | |
] | |
generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input]) | |
if __name__ == "__main__": | |
demo.queue().launch(mcp_server=True) |