import torch from typing import List from diffusers import AutoencoderKLWan, WanVACEPipeline, UniPCMultistepScheduler from diffusers.utils import export_to_video 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 from briarmbg import BriaRMBG model_id = "Wan-AI/Wan2.1-VACE-14B-diffusers" vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) pipe = WanVACEPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16).to("cuda") # Initialize background removal model rmbg = BriaRMBG.from_pretrained("briaai/RMBG-1.4").to("cuda", dtype=torch.float32) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0) 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 = 16 MIN_FRAMES_MODEL = 8 MAX_FRAMES_MODEL = 81 # Default prompts for different modes - Updated with new mode names MODE_PROMPTS = { "Reference": "the playful penguin picks up the green cat eye sunglasses and puts them on", "First - Last Frame": "CG animation style, a small blue bird takes off from the ground, flapping its wings. The bird's feathers are delicate, with a unique pattern on its chest. The background shows a blue sky with white clouds under bright sunshine. The camera follows the bird upward, capturing its flight and the vastness of the sky from a close-up, low-angle perspective.", "Random Transitions": "Various different characters appear and disappear in a fast transition video showcasting their unique features and personalities. The video is about showcasing different dance styles, with each character performing a distinct dance move. The background is a vibrant, colorful stage with dynamic lighting that changes with each dance style. The camera captures close-ups of the dancers' expressions and movements. Highly dynamic, fast-paced music video, with quick cuts and transitions between characters, cinematic, vibrant colors" } 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 remove_alpha_channel(image: Image.Image) -> Image.Image: """ Remove alpha channel from PIL Image if it exists. Args: image (Image.Image): Input PIL image Returns: Image.Image: Image with alpha channel removed (RGB format) """ if image.mode in ('RGBA', 'LA'): # Create a white background background = Image.new('RGB', image.size, (255, 255, 255)) # Paste the image onto the white background using alpha channel as mask if image.mode == 'RGBA': background.paste(image, mask=image.split()[-1]) # Use alpha channel as mask else: # LA mode background.paste(image.convert('RGB'), mask=image.split()[-1]) return background elif image.mode == 'P': # Convert palette mode to RGB (some palette images have transparency) if 'transparency' in image.info: image = image.convert('RGBA') background = Image.new('RGB', image.size, (255, 255, 255)) background.paste(image, mask=image.split()[-1]) return background else: return image.convert('RGB') elif image.mode != 'RGB': # Convert any other mode to RGB return image.convert('RGB') else: # Already RGB, return as is return image @torch.inference_mode() def numpy2pytorch(imgs): h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0 # so that 127 must be strictly 0.0 h = h.movedim(-1, 1) return h @torch.inference_mode() def pytorch2numpy(imgs, quant=True): results = [] for x in imgs: y = x.movedim(0, -1) if quant: y = y * 127.5 + 127.5 y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8) else: y = y * 0.5 + 0.5 y = y.detach().float().cpu().numpy().clip(0, 1).astype(np.float32) results.append(y) return results def resize_without_crop(image, target_width, target_height): pil_image = Image.fromarray(image) resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS) return np.array(resized_image) @torch.inference_mode() def run_rmbg(img, sigma=0.0): """ Remove background from image using BriaRMBG model. Args: img (np.ndarray): Input image as numpy array (H, W, C) sigma (float): Noise parameter for blending Returns: tuple: (result_image, alpha_mask) where result_image is the image with background removed """ H, W, C = img.shape assert C == 3 k = (256.0 / float(H * W)) ** 0.5 feed = resize_without_crop(img, int(64 * round(W * k)), int(64 * round(H * k))) feed = numpy2pytorch([feed]).to(device="cuda", dtype=torch.float32) alpha = rmbg(feed)[0][0] alpha = torch.nn.functional.interpolate(alpha, size=(H, W), mode="bilinear") alpha = alpha.movedim(1, -1)[0] alpha = alpha.detach().float().cpu().numpy().clip(0, 1) result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha return result.clip(0, 255).astype(np.uint8), alpha def remove_background_from_image(image: Image.Image) -> Image.Image: """ Remove background from PIL Image using RMBG model. Args: image (Image.Image): Input PIL image Returns: Image.Image: Image with background removed (transparent background) """ # Convert PIL to numpy array img_array = np.array(image) # Remove background using RMBG result_array, alpha_mask = run_rmbg(img_array) # Convert back to PIL with alpha channel result_image = Image.fromarray(result_array) # Create RGBA image with alpha mask if result_image.mode != 'RGBA': result_image = result_image.convert('RGBA') # Handle alpha mask dimensions and convert to PIL # The alpha_mask might have extra dimensions, so squeeze and ensure 2D alpha_mask_2d = np.squeeze(alpha_mask) if alpha_mask_2d.ndim > 2: # If still more than 2D, take the first channel alpha_mask_2d = alpha_mask_2d[:, :, 0] if alpha_mask_2d.shape[-1] == 1 else alpha_mask_2d[:, :, 0] # Convert to uint8 and create PIL Image without deprecated mode parameter alpha_array = (alpha_mask_2d * 255).astype(np.uint8) alpha_pil = Image.fromarray(alpha_array, 'L') result_image.putalpha(alpha_pil) return result_image 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][0] # Remove alpha channel before calculating dimensions first_image = remove_alpha_channel(first_image) 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 prepare_video_and_mask_Ref2V(height: int, width: int, num_frames: int): frames = [] # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to # match the original code. frames.extend([Image.new("RGB", (width, height), (128, 128, 128))] * (num_frames)) mask_white = Image.new("L", (width, height), 255) mask = [mask_white] * (num_frames) return frames, mask def prepare_video_and_mask_FLF2V(first_img: Image.Image, last_img: Image.Image, height: int, width: int, num_frames: int): # Remove alpha channels before processing first_img = remove_alpha_channel(first_img) last_img = remove_alpha_channel(last_img) first_img = first_img.resize((width, height)) last_img = last_img.resize((width, height)) frames = [] frames.append(first_img) # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to # match the original code. frames.extend([Image.new("RGB", (width, height), (128, 128, 128))] * (num_frames - 2)) frames.append(last_img) mask_black = Image.new("L", (width, height), 0) mask_white = Image.new("L", (width, height), 255) mask = [mask_black, *[mask_white] * (num_frames - 2), mask_black] return frames, mask def calculate_random2v_frame_indices(num_images: int, num_frames: int) -> List[int]: """ Calculate evenly spaced frame indices for Random2V mode. Args: num_images (int): Number of input images num_frames (int): Total number of frames in the video Returns: List[int]: Frame indices where images should be placed """ if num_images <= 0: return [] if num_images == 1: # Single image goes in the middle return [num_frames // 2] if num_images >= num_frames: # More images than frames, use every frame return list(range(num_frames)) # Calculate evenly spaced indices # We want to distribute images across the full duration indices = [] step = (num_frames - 1) / (num_images - 1) for i in range(num_images): frame_idx = int(round(i * step)) # Ensure we don't exceed num_frames - 1 frame_idx = min(frame_idx, num_frames - 1) indices.append(frame_idx) # Remove duplicates while preserving order seen = set() unique_indices = [] for idx in indices: if idx not in seen: seen.add(idx) unique_indices.append(idx) return unique_indices def prepare_video_and_mask_Random2V(images: List[Image.Image], frame_indices: List[int], height: int, width: int, num_frames: int): # Remove alpha channels from all images before processing images = [remove_alpha_channel(img) for img in images] images = [img.resize((width, height)) for img in images] # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to # match the original code. frames = [Image.new("RGB", (width, height), (128, 128, 128))] * num_frames mask_black = Image.new("L", (width, height), 0) mask_white = Image.new("L", (width, height), 255) mask = [mask_white] * num_frames for img, idx in zip(images, frame_indices): assert idx < num_frames, f"Frame index {idx} exceeds num_frames {num_frames}" frames[idx] = img mask[idx] = mask_black return frames, mask def get_duration(gallery_images, mode, prompt, height, width, negative_prompt, duration_seconds, guidance_scale, steps, seed, randomize_seed, remove_bg, progress): # Add extra time if background removal is enabled base_duration = 60 if steps > 4 and duration_seconds > 2: base_duration = 90 elif steps > 4 or duration_seconds > 2: base_duration = 75 # Add extra time for background removal processing if mode == "Reference" and remove_bg: # Updated to use new mode name base_duration += 30 return base_duration @spaces.GPU(duration=get_duration) 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, remove_bg = 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 - "Reference", "first - last frame", or "random transitions" 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 remove_bg (bool): Whether to remove background from images (reference mode only) 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.") else: # Process images: remove background if requested (reference mode only), then remove alpha channels processed_images = [] for img in gallery_images: image = img[0] # Extract PIL image from gallery format # Apply background removal only for reference mode if checkbox is checked if mode == "Reference" and remove_bg: # Updated to use new mode name image = remove_background_from_image(image) # Always remove alpha channels to ensure RGB format image = remove_alpha_channel(image) processed_images.append(image) gallery_images = processed_images if mode == "First - Last Frame" and len(gallery_images) >= 2: # Updated mode name gallery_images = gallery_images[:2] elif mode == "First - Last Frame" and len(gallery_images) < 2: # Updated mode name raise gr.Error("First - Last Frame mode requires at least 2 images, but only {} were supplied.".format(len(gallery_images))) 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) # Process images based on the selected mode if mode == "First - Last Frame": # Updated mode name frames, mask = prepare_video_and_mask_FLF2V( first_img=gallery_images[0], last_img=gallery_images[1], height=target_h, width=target_w, num_frames=num_frames ) reference_images = None elif mode == "Reference": # Updated mode name frames, mask = prepare_video_and_mask_Ref2V(height=target_h, width=target_w, num_frames=num_frames) reference_images = gallery_images else: # mode == "random transitions" # Updated mode name # Calculate dynamic frame indices based on number of images and frames frame_indices = calculate_random2v_frame_indices(len(gallery_images), num_frames) frames, mask = prepare_video_and_mask_Random2V( images=gallery_images, frame_indices=frame_indices, height=target_h, width=target_w, num_frames=num_frames ) reference_images = None with torch.inference_mode(): output_frames_list = pipe( video=frames, mask=mask, reference_images=reference_images, 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 control_modes = """ **3 Control Modes Available:** - **Reference** Generate a video incorporating elements from input reference images - **First - Last Frame** Generate a video using first and last frame conditioning defined by input images - **Random Transitions** Generate a video with intermediate transitions between multiple input images """ with gr.Blocks() as demo: gr.Markdown("# Fast 6 step Wan 2.1 VACE (14B)") gr.Markdown("Using [**Wan2.1-VACE-14B**](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) + [**👻FusionX Phantom LoRA**](https://huggingface.co/vrgamedevgirl84/Wan14BT2VFusioniX) by [**vrgamedevgirl84**](https://huggingface.co/vrgamedevgirl84) with **🧨diffusers**, for fast video generation with multiple conditions 🏎️") gr.Markdown(f"{control_modes}") with gr.Row(): with gr.Column(): with gr.Group(): # Radio button for mode selection with updated names mode_radio = gr.Radio( choices=["Reference", "First - Last Frame", "Random Transitions"], value="reference", label="Control Mode", #info="Reference: upload reference images to take elements from | First - Last Frame: upload 1st and last frames| Random Transitions: upload images to be used as frame anchors" ) # Gallery component for multiple image upload gallery_component = gr.Gallery( label="upload 1 or more images", show_label=True, elem_id="gallery", columns=3, rows=2, object_fit="contain", height="auto", type="pil", allow_preview=True ) # Background removal checkbox moved here - right beneath control modes remove_bg_checkbox = gr.Checkbox( label="Remove Background", value=False, info="removes background from input images, enable to prevent unwanted background elements in the generated video" ) prompt_input = gr.Textbox(label="Prompt", value=MODE_PROMPTS["Reference"]) 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.3, 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=10, 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) # Function to update checkbox visibility based on mode def update_bg_removal_visibility(mode): return gr.update(visible=(mode == "reference")) # Updated to use new mode name # Update prompt when mode changes mode_radio.change( fn=update_prompt_from_mode, inputs=[mode_radio], outputs=[prompt_input] ) # Update background removal checkbox visibility when mode changes mode_radio.change( fn=update_bg_removal_visibility, inputs=[mode_radio], outputs=[remove_bg_checkbox] ) # 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, remove_bg_checkbox ] generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input]) gr.Examples( examples=[ [["reachy.png", "sunglasses.jpg", "gpu_hat.png"], "reference", "the cute robot is wearing the sunglasses and the hat that reads 'GPU poor', and moves around playfully", 480, 832], [["flf2v_input_first_frame.png", "flf2v_input_last_frame.png"], "first - last frame", "CG animation style, a small blue bird takes off from the ground, flapping its wings. The bird's feathers are delicate, with a unique pattern on its chest. The background shows a blue sky with white clouds under bright sunshine. The camera follows the bird upward, capturing its flight and the vastness of the sky from a close-up, low-angle perspective.", 512, 512], ], inputs=[gallery_component, mode_radio, prompt_input, height_input, width_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy" ) if __name__ == "__main__": demo.queue().launch(mcp_server=True)