import os import spaces import shutil import subprocess import sys import copy import random import tempfile import warnings import time import gc import uuid from tqdm import tqdm import cv2 import numpy as np import torch from torch.nn import functional as F from PIL import Image import gradio as gr from diffusers import ( FlowMatchEulerDiscreteScheduler, SASolverScheduler, DEISMultistepScheduler, DPMSolverMultistepInverseScheduler, UniPCMultistepScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, ) from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline from diffusers.utils.export_utils import export_to_video from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig import aoti os.environ["TOKENIZERS_PARALLELISM"] = "true" warnings.filterwarnings("ignore") IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU")) # --- FIX: Use ZeroGPU ephemeral storage --- if IS_ZERO_GPU: CACHE_DIR = "/data-nvme/huggingface_cache" else: CACHE_DIR = os.path.expanduser("~/.cache/huggingface/") os.environ["HF_HOME"] = CACHE_DIR os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR os.environ["DIFFUSERS_CACHE"] = CACHE_DIR os.environ["HUGGINGFACE_HUB_CACHE"] = CACHE_DIR # --- FIX: Download ALL RIFE model files needed --- def setup_rife_complete(): if not os.path.exists("RIFEv4.26_0921.zip"): print("Downloading RIFE Model...") subprocess.run([ "wget", "-q", "https://huggingface.co/r3gm/RIFE/resolve/main/RIFEv4.26_0921.zip", "-O", "RIFEv4.26_0921.zip" ], check=True) subprocess.run(["unzip", "-o", "RIFEv4.26_0921.zip"], check=True) # FIX: Download ALL model files from RIFE repo (not just warplayer.py) if not os.path.exists("model"): os.makedirs("model", exist_ok=True) # List of ALL files needed from model/ folder model_files = [ "warplayer.py", "IFNet.py", "RIFE.py", "refine.py", "loss.py", "IFNet_2R.py", "IFNet_m.py", "IFNet_hdv3.py", "RIFE_HDv3.py" ] base_url = "https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/main/model/" for file in model_files: url = base_url + file output_path = f"model/{file}" if not os.path.exists(output_path): print(f"Downloading {file}...") result = subprocess.run( ["wget", "-q", url, "-O", output_path], capture_output=True, text=True ) if result.returncode != 0: print(f"Warning: Could not download {file}, trying alternative...") # Try alternative URL patterns alt_urls = [ f"https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/master/model/{file}", f"https://raw.githubusercontent.com/hzwer/RIFE/main/model/{file}", ] for alt_url in alt_urls: result = subprocess.run( ["wget", "-q", alt_url, "-O", output_path], capture_output=True, text=True ) if result.returncode == 0: break # Create __init__.py with open("model/__init__.py", "w") as f: f.write("") print("RIFE model/ folder setup complete") setup_rife_complete() sys.path.insert(0, os.getcwd()) from train_log.RIFE_HDv3 import Model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") rife_model = Model() rife_model.load_model("train_log", -1) rife_model.eval() # --- FRAME EXTRACTION --- get_timestamp_js = """ function() { const video = document.querySelector('#generated-video video'); if (video) { return video.currentTime; } else { return 0; } } """ def extract_frame(video_path, timestamp): if not video_path: return None cap = cv2.VideoCapture(video_path) if not cap.isOpened(): return None fps = cap.get(cv2.CAP_PROP_FPS) target_frame_num = int(float(timestamp) * fps) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if target_frame_num >= total_frames: target_frame_num = total_frames - 1 cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame_num) ret, frame = cap.read() cap.release() if ret: return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) return None def clear_vram(): gc.collect() torch.cuda.empty_cache() @torch.no_grad() def interpolate_bits(frames_np, multiplier=2, scale=1.0): if isinstance(frames_np, list): T = len(frames_np) H, W, C = frames_np[0].shape else: T, H, W, C = frames_np.shape if multiplier < 2: if isinstance(frames_np, np.ndarray): return list(frames_np) return frames_np n_interp = multiplier - 1 tmp = max(128, int(128 / scale)) ph = ((H - 1) // tmp + 1) * tmp pw = ((W - 1) // tmp + 1) * tmp padding = (0, pw - W, 0, ph - H) def to_tensor(frame_np): t = torch.from_numpy(frame_np).to(device) t = t.permute(2, 0, 1).unsqueeze(0) return F.pad(t, padding).half() def from_tensor(tensor): t = tensor[0, :, :H, :W] t = t.permute(1, 2, 0) return t.float().cpu().numpy() def make_inference(I0, I1, n): if rife_model.version >= 3.9: res = [] for i in range(n): res.append(rife_model.inference(I0, I1, (i+1) * 1. / (n+1), scale)) return res else: middle = rife_model.inference(I0, I1, scale) if n == 1: return [middle] first_half = make_inference(I0, middle, n=n//2) second_half = make_inference(middle, I1, n=n//2) if n % 2: return [*first_half, middle, *second_half] else: return [*first_half, *second_half] output_frames = [] I1 = to_tensor(frames_np[0]) total_steps = T - 1 with tqdm(total=total_steps, desc="Interpolating", unit="frame") as pbar: for i in range(total_steps): I0 = I1 output_frames.append(from_tensor(I0)) I1 = to_tensor(frames_np[i+1]) mid_tensors = make_inference(I0, I1, n_interp) for mid in mid_tensors: output_frames.append(from_tensor(mid)) if (i + 1) % 50 == 0: pbar.update(50) pbar.update(total_steps % 50) output_frames.append(from_tensor(I1)) del I0, I1, mid_tensors torch.cuda.empty_cache() return output_frames # --- WAN PIPELINE --- # FIX: Use 1.3B model to stay under 50GB storage limit MODEL_ID = "Wan-AI/Wan2.2-I2V-T2V-1.3B-Diffusers" MAX_DIM = 832 MIN_DIM = 480 SQUARE_DIM = 640 MULTIPLE_OF = 16 MAX_SEED = np.iinfo(np.int32).max FIXED_FPS = 16 MIN_FRAMES_MODEL = 8 MAX_FRAMES_MODEL = 160 MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1) MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1) SCHEDULER_MAP = { "FlowMatchEulerDiscrete": FlowMatchEulerDiscreteScheduler, "SASolver": SASolverScheduler, "DEISMultistep": DEISMultistepScheduler, "DPMSolverMultistepInverse": DPMSolverMultistepInverseScheduler, "UniPCMultistep": UniPCMultistepScheduler, "DPMSolverMultistep": DPMSolverMultistepScheduler, "DPMSolverSinglestep": DPMSolverSinglestepScheduler, } print(f"Loading model from {MODEL_ID}...") pipe = WanImageToVideoPipeline.from_pretrained( MODEL_ID, torch_dtype=torch.bfloat16, cache_dir=CACHE_DIR, ).to('cuda') original_scheduler = copy.deepcopy(pipe.scheduler) # Quantize to save VRAM quantize_(pipe.text_encoder, Int8WeightOnlyConfig()) quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig()) if hasattr(pipe, 'transformer_2') and pipe.transformer_2 is not None: quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig()) # Only use aoti if available try: aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da') if hasattr(pipe, 'transformer_2') and pipe.transformer_2 is not None: aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da') except Exception as e: print(f"AoT compilation not available: {e}") default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation" default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走" def resize_image(image: Image.Image) -> Image.Image: width, height = image.size if width == height: return image.resize((SQUARE_DIM, SQUARE_DIM), Image.LANCZOS) aspect_ratio = width / height MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM image_to_resize = image if aspect_ratio > MAX_ASPECT_RATIO: target_w, target_h = MAX_DIM, MIN_DIM crop_width = int(round(height * MAX_ASPECT_RATIO)) left = (width - crop_width) // 2 image_to_resize = image.crop((left, 0, left + crop_width, height)) elif aspect_ratio < MIN_ASPECT_RATIO: target_w, target_h = MIN_DIM, MAX_DIM crop_height = int(round(width / MIN_ASPECT_RATIO)) top = (height - crop_height) // 2 image_to_resize = image.crop((0, top, width, top + crop_height)) else: if width > height: target_w = MAX_DIM target_h = int(round(target_w / aspect_ratio)) else: target_h = MAX_DIM target_w = int(round(target_h * aspect_ratio)) final_w = round(target_w / MULTIPLE_OF) * MULTIPLE_OF final_h = round(target_h / MULTIPLE_OF) * MULTIPLE_OF final_w = max(MIN_DIM, min(MAX_DIM, final_w)) final_h = max(MIN_DIM, min(MAX_DIM, final_h)) return image_to_resize.resize((final_w, final_h), Image.LANCZOS) def resize_and_crop_to_match(target_image, reference_image): ref_width, ref_height = reference_image.size target_width, target_height = target_image.size scale = max(ref_width / target_width, ref_height / target_height) new_width, new_height = int(target_width * scale), int(target_height * scale) resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS) left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2 return resized.crop((left, top, left + ref_width, top + ref_height)) def get_num_frames(duration_seconds: float): return 1 + int(np.clip( int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL, )) def get_inference_duration( resized_image, processed_last_image, prompt, steps, negative_prompt, num_frames, guidance_scale, guidance_scale_2, current_seed, scheduler_name, flow_shift, frame_multiplier, quality, duration_seconds, progress ): BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624 BASE_STEP_DURATION = 15 width, height = resized_image.size factor = num_frames * width * height / BASE_FRAMES_HEIGHT_WIDTH step_duration = BASE_STEP_DURATION * factor ** 1.5 gen_time = int(steps) * step_duration if guidance_scale > 1: gen_time = gen_time * 1.8 frame_factor = frame_multiplier // FIXED_FPS if frame_factor > 1: total_out_frames = (num_frames * frame_factor) - num_frames inter_time = (total_out_frames * 0.02) gen_time += inter_time return 15 + gen_time @spaces.GPU(duration=get_inference_duration) def run_inference( resized_image, processed_last_image, prompt, steps, negative_prompt, num_frames, guidance_scale, guidance_scale_2, current_seed, scheduler_name, flow_shift, frame_multiplier, quality, duration_seconds, progress=gr.Progress(track_tqdm=True), ): scheduler_class = SCHEDULER_MAP.get(scheduler_name) if scheduler_class.__name__ != pipe.scheduler.config._class_name or flow_shift != pipe.scheduler.config.get("flow_shift", "shift"): config = copy.deepcopy(original_scheduler.config) if scheduler_class == FlowMatchEulerDiscreteScheduler: config['shift'] = flow_shift else: config['flow_shift'] = flow_shift pipe.scheduler = scheduler_class.from_config(config) clear_vram() task_name = str(uuid.uuid4())[:8] print(f"Task: {task_name}, {duration_seconds}, {resized_image.size}, FM={frame_multiplier}") result = pipe( image=resized_image, last_image=processed_last_image, prompt=prompt, negative_prompt=negative_prompt, height=resized_image.height, width=resized_image.width, num_frames=num_frames, guidance_scale=float(guidance_scale), guidance_scale_2=float(guidance_scale_2), num_inference_steps=int(steps), generator=torch.Generator(device="cuda").manual_seed(current_seed), output_type="np" ) raw_frames_np = result.frames[0] pipe.scheduler = original_scheduler frame_factor = frame_multiplier // FIXED_FPS if frame_factor > 1: rife_model.device() rife_model.flownet = rife_model.flownet.half() final_frames = interpolate_bits(raw_frames_np, multiplier=int(frame_factor)) else: final_frames = list(raw_frames_np) final_fps = FIXED_FPS * int(frame_factor) with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile: video_path = tmpfile.name with tqdm(total=3, desc="Rendering Media", unit="clip") as pbar: pbar.update(2) export_to_video(final_frames, video_path, fps=final_fps, quality=quality) pbar.update(1) del result, raw_frames_np, final_frames clear_vram() return video_path, task_name def generate_video( input_image, last_image, prompt, steps=4, negative_prompt=default_negative_prompt, duration_seconds=MAX_DURATION, guidance_scale=1, guidance_scale_2=1, seed=42, randomize_seed=False, quality=5, scheduler="UniPCMultistep", flow_shift=6.0, frame_multiplier=16, video_component=True, progress=gr.Progress(track_tqdm=True), ): if input_image is None: raise gr.Error("Please upload an input image.") num_frames = get_num_frames(duration_seconds) current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) resized_image = resize_image(input_image) processed_last_image = None if last_image: processed_last_image = resize_and_crop_to_match(last_image, resized_image) video_path, task_n = run_inference( resized_image, processed_last_image, prompt, steps, negative_prompt, num_frames, guidance_scale, guidance_scale_2, current_seed, scheduler, flow_shift, frame_multiplier, quality, duration_seconds, progress, ) print(f"GPU complete: {task_n}") return (video_path if video_component else None), video_path, current_seed CSS = """ #hidden-timestamp { opacity: 0; height: 0px; width: 0px; margin: 0px; padding: 0px; overflow: hidden; position: absolute; pointer-events: none; } """ with gr.Blocks(theme=gr.themes.Soft(), css=CSS, delete_cache=(3600, 10800)) as demo: gr.Markdown("## WAMU V2 - Wan 2.2 I2V (1.3B) 🐢") gr.Markdown("#### ℹ️ **Lightweight Version:** Uses 1.3B model for Spaces compatibility.") gr.Markdown("Run Wan 2.2 in just 4-8 steps, fp8 quantization - compatible with 🧨 diffusers and ZeroGPU.") with gr.Row(): with gr.Column(): input_image_component = gr.Image(type="pil", label="Input Image", sources=["upload", "clipboard"]) prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v) duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=3.5, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.") frame_multi = gr.Dropdown( choices=[FIXED_FPS, FIXED_FPS*2, FIXED_FPS*4], value=FIXED_FPS, label="Video Fluidity (Frames per Second)", info="Extra frames will be generated using flow estimation." ) with gr.Accordion("Advanced Settings", open=False): last_image_component = gr.Image(type="pil", label="Last Image (Optional)", sources=["upload", "clipboard"]) negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, info="Used if any Guidance Scale > 1.", lines=3) quality_slider = gr.Slider(minimum=1, maximum=10, step=1, value=6, label="Video Quality") 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) steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps") guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale") guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2") scheduler_dropdown = gr.Dropdown( label="Scheduler", choices=list(SCHEDULER_MAP.keys()), value="UniPCMultistep", info="Select a custom scheduler." ) flow_shift_slider = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=3.0, label="Flow Shift") play_result_video = gr.Checkbox(label="Display result", value=True, interactive=True) generate_button = gr.Button("Generate Video", variant="primary") with gr.Column(): video_output = gr.Video(label="Generated Video", autoplay=True, sources=["upload"], show_download_button=True, show_share_button=True, interactive=False, elem_id="generated-video") with gr.Row(): grab_frame_btn = gr.Button("📸 Use Current Frame as Input", variant="secondary") timestamp_box = gr.Number(value=0, label="Timestamp", visible=True, elem_id="hidden-timestamp") file_output = gr.File(label="Download Video") ui_inputs = [ input_image_component, last_image_component, prompt_input, steps_slider, negative_prompt_input, duration_seconds_input, guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox, quality_slider, scheduler_dropdown, flow_shift_slider, frame_multi, play_result_video ] generate_button.click( fn=generate_video, inputs=ui_inputs, outputs=[video_output, file_output, seed_input] ) grab_frame_btn.click( fn=None, inputs=None, outputs=[timestamp_box], js=get_timestamp_js ) timestamp_box.change( fn=extract_frame, inputs=[video_output, timestamp_box], outputs=[input_image_component] ) if __name__ == "__main__": demo.queue().launch(mcp_server=True, ssr_mode=False, show_error=True)