import spaces import os import gc import cv2 import json import uuid import shutil import subprocess import tempfile import threading import zipfile from pathlib import Path import torch print("Torch:", torch.__version__) print("Torch CUDA:", torch.version.cuda) print("CUDA available:", torch.cuda.is_available()) import numpy as np import gradio as gr from PIL import Image from huggingface_hub import snapshot_download, hf_hub_download # ============================================================ # DIFFUSERS # ============================================================ import diffusers print("Diffusers version:", diffusers.__version__) print( "Has WanVACEPipeline:", hasattr(diffusers, "WanVACEPipeline") ) from diffusers import ( WanVACEPipeline, WanVACETransformer3DModel, AutoencoderKLWan, ) from diffusers.schedulers.scheduling_unipc_multistep import ( UniPCMultistepScheduler, ) from diffusers.utils import export_to_video # ============================================================ # SAM 2.1 video mask propagation # ============================================================ try: from sam2.build_sam import build_sam2_video_predictor except Exception as _sam2_import_error: build_sam2_video_predictor = None _SAM2_IMPORT_ERROR = _sam2_import_error else: _SAM2_IMPORT_ERROR = None # ============================================================ # CONFIG # ============================================================ WAN_MODEL = os.environ.get( "WAN_MODEL", "Wan-AI/Wan2.1-VACE-14B-diffusers", ) # SAM 2.1 propagates a mask inferred from the original-vs-edited first-frame difference. SAM2_CONFIG = os.environ.get("SAM2_CONFIG", "configs/sam2.1/sam2.1_hiera_l.yaml") SAM2_CHECKPOINT_URL = os.environ.get( "SAM2_CHECKPOINT_URL", "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt", ) # Prefer persistent HF Space storage when it exists; otherwise fall back # to /tmp for a single-runtime cache. _DEFAULT_MODEL_ROOT = ( Path("/data/zero_drift_models") if Path("/data").exists() else Path("/tmp/zero_drift_models") ) MODEL_ROOT = Path( os.environ.get( "MODEL_ROOT", str(_DEFAULT_MODEL_ROOT), ) ) WAN_LOCAL_DIR = MODEL_ROOT / "wan_vace" SAM2_LOCAL_DIR = MODEL_ROOT / "sam2" SAM2_CHECKPOINT = SAM2_LOCAL_DIR / "sam2.1_hiera_large.pt" # ZeroGPU exposes CUDA when the @spaces.GPU function is executing. DEVICE = "cuda" WAN_DTYPE = torch.bfloat16 # Favor Tensor Core execution for the remaining FP32 helper operations. try: torch.set_float32_matmul_precision("high") except Exception: pass try: torch.backends.cuda.matmul.allow_tf32 = True except Exception: pass try: torch.backends.cudnn.allow_tf32 = True except Exception: pass # Full-quality 14B BF16 path. The 96 GB xlarge GPU has enough VRAM, # so do not quantize the transformer: native BF16 Tensor Core execution # is the latency-first path here. USE_TORCH_COMPILE = False # Stable eager path; avoids the current Dynamo/CUDA handler bug. OUTPUT_DIR = Path( "/tmp/zero_drift_outputs" ) OUTPUT_DIR.mkdir( parents=True, exist_ok=True, ) # Maximum source duration. MAX_SECONDS = 15 # Wan 2.1 VACE is trained around short 4k+1-frame clips. # Longer source videos are processed in overlapping chunks. VACE_CHUNK_FRAMES = 81 VACE_CHUNK_OVERLAP = 8 MODEL_FPS = 16.0 # 14B quality path: keep the full 720p-capable spatial budget. # Speed is prioritized with native BF16 execution and low step counts. MAX_WIDTH = 1280 MAX_HEIGHT = 720 # ============================================================ # GLOBAL MODELS # ============================================================ WAN_PIPE = None SAM2_PREDICTOR = None MODEL_LOCK = threading.Lock() # ============================================================ # STARTUP MODEL DOWNLOADS # ============================================================ def ensure_model_downloads(): """Download the Wan VACE repository and public SAM 2.1 checkpoint if absent.""" MODEL_ROOT.mkdir(parents=True, exist_ok=True) WAN_LOCAL_DIR.mkdir(parents=True, exist_ok=True) SAM2_LOCAL_DIR.mkdir(parents=True, exist_ok=True) print("=" * 72) print("ZERO-DRIFT STARTUP MODEL DOWNLOAD") print("Model root:", MODEL_ROOT) print("Wan repo:", WAN_MODEL) print("SAM 2.1 checkpoint:", SAM2_CHECKPOINT) print("=" * 72) wan_config = WAN_LOCAL_DIR / "model_index.json" wan_vae_config = WAN_LOCAL_DIR / "vae" / "config.json" wan_vae_weights = WAN_LOCAL_DIR / "vae" / "diffusion_pytorch_model.safetensors" wan_cache_valid = wan_config.exists() and wan_vae_config.exists() and wan_vae_weights.exists() if wan_cache_valid: print("Wan 2.1 VACE Diffusers repo already present:", WAN_LOCAL_DIR) else: if WAN_LOCAL_DIR.exists(): print("Wan cache is incomplete. Removing stale cache:", WAN_LOCAL_DIR) shutil.rmtree(WAN_LOCAL_DIR, ignore_errors=True) WAN_LOCAL_DIR.mkdir(parents=True, exist_ok=True) print("Downloading Wan 2.1 VACE Diffusers. This repository is about 75 GB.") snapshot_download(repo_id=WAN_MODEL, local_dir=str(WAN_LOCAL_DIR)) if not wan_vae_config.exists() or not wan_vae_weights.exists(): raise RuntimeError("Wan download finished but required VAE files are missing.") print("Wan 2.1 VACE download complete.") if not SAM2_CHECKPOINT.exists() or SAM2_CHECKPOINT.stat().st_size == 0: print("Downloading public SAM 2.1 Hiera-L checkpoint...") import urllib.request partial = SAM2_CHECKPOINT.with_suffix(".pt.partial") urllib.request.urlretrieve(SAM2_CHECKPOINT_URL, str(partial)) partial.replace(SAM2_CHECKPOINT) print("SAM 2.1 checkpoint ready:", SAM2_CHECKPOINT) print("=" * 72) # ============================================================ # GENERAL HELPERS # ============================================================ def cleanup(): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() try: torch.cuda.ipc_collect() except Exception: pass def run_command(command): process = subprocess.run( command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) if process.returncode != 0: raise RuntimeError( "\n".join([ "Command failed:", " ".join(command), "", process.stderr[-6000:], ]) ) return process.stdout def get_video_path(value): if value is None: return None if isinstance(value, str): return value if isinstance(value, dict): if value.get("path"): return value["path"] if value.get("video"): return value["video"] raise ValueError( f"Unsupported video input: {type(value)}" ) # ============================================================ # VIDEO PROBING # ============================================================ def probe_video(path): result = run_command([ "ffprobe", "-v", "error", "-show_entries", "stream=index,codec_type,width,height," "r_frame_rate,avg_frame_rate,duration", "-of", "json", path, ]) data = json.loads(result) video_stream = None audio_stream = None for stream in data.get("streams", []): if stream.get("codec_type") == "video": video_stream = stream elif stream.get("codec_type") == "audio": audio_stream = stream if video_stream is None: raise RuntimeError( "Input does not contain a video stream." ) width = int( video_stream["width"] ) height = int( video_stream["height"] ) fps_string = ( video_stream.get("avg_frame_rate") or video_stream.get("r_frame_rate") or "16/1" ) try: n, d = fps_string.split("/") fps = float(n) / float(d) except Exception: fps = 16.0 duration = float( video_stream.get( "duration", 0, ) or 0 ) return { "width": width, "height": height, "fps": fps, "duration": duration, "has_audio": audio_stream is not None, } # ============================================================ # AUDIO # ============================================================ def extract_audio( input_video, output_audio, ): # -vn means absolutely no video processing. # -c:a copy means the original audio stream is copied. run_command([ "ffmpeg", "-y", "-i", input_video, "-vn", "-map", "0:a:0?", "-c:a", "copy", output_audio, ]) return output_audio # ============================================================ # VIDEO DECODING # ============================================================ def decode_frames( video_path, max_frames=None, target_fps=MODEL_FPS, ): info = probe_video( video_path ) cap = cv2.VideoCapture( video_path ) if not cap.isOpened(): raise RuntimeError( "Could not open input video." ) source_fps = float( info.get("fps") or target_fps ) duration = float( info.get("duration") or 0.0 ) if target_fps and duration > 0: target_count = max(1, int(round(duration * target_fps))) else: target_count = int( cap.get(cv2.CAP_PROP_FRAME_COUNT) ) source_count = max( 1, int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), ) if max_frames is not None: target_count = min( target_count, int(max_frames), ) target_count = min( target_count, source_count, ) target_indices = np.linspace( 0, max(0, source_count - 1), target_count, dtype=np.int64, ) target_set = { int(i): pos for pos, i in enumerate(target_indices) } frames = [] source_index = 0 while len(frames) < target_count: ok, frame = cap.read() if not ok: break if source_index in target_set: frame = cv2.cvtColor( frame, cv2.COLOR_BGR2RGB, ) frames.append( Image.fromarray(frame) ) source_index += 1 cap.release() if not frames: raise RuntimeError( "Could not decode any frames." ) # The Wan VACE model is used at a stable 16 fps working rate. # This preserves the source duration while avoiding unnecessary # extra VACE chunks for 24/30/60 fps inputs. return ( frames, float(target_fps or source_fps), info, ) # ============================================================ # RESOLUTION # ============================================================ def resize_for_wan( frames, ): width, height = frames[0].size scale = min( MAX_WIDTH / width, MAX_HEIGHT / height, 1.0, ) new_width = int( width * scale ) new_height = int( height * scale ) # Wan dimensions should be divisible by 8. new_width -= ( new_width % 8 ) new_height -= ( new_height % 8 ) new_width = max( 256, new_width, ) new_height = max( 256, new_height, ) if ( new_width == width and new_height == height ): return frames return [ frame.resize( ( new_width, new_height, ), Image.Resampling.LANCZOS, ) for frame in frames ] # ============================================================ # MODEL LOADING # ============================================================ def load_models(): global WAN_PIPE, SAM2_PREDICTOR with MODEL_LOCK: if WAN_PIPE is None: print("Loading Wan 2.1 VACE Diffusers from local files...", WAN_LOCAL_DIR) vae_config = WAN_LOCAL_DIR / "vae" / "config.json" vae_weights = WAN_LOCAL_DIR / "vae" / "diffusion_pytorch_model.safetensors" if not vae_config.exists() or not vae_weights.exists(): raise RuntimeError(f"Wan VACE cache is incomplete: expected {vae_config} and {vae_weights}.") vae = AutoencoderKLWan.from_pretrained( str(WAN_LOCAL_DIR), subfolder="vae", torch_dtype=torch.float32, local_files_only=True ) WAN_PIPE = WanVACEPipeline.from_pretrained( str(WAN_LOCAL_DIR), vae=vae, torch_dtype=WAN_DTYPE, local_files_only=True ) WAN_PIPE.scheduler = UniPCMultistepScheduler.from_config(WAN_PIPE.scheduler.config, flow_shift=5.0) WAN_PIPE.to(DEVICE) try: WAN_PIPE.fuse_qkv_projections() print("VACE QKV projections fused.") except Exception as e: print("QKV fusion skipped:", repr(e)) print("torch.compile disabled; using native BF16 eager execution.") WAN_PIPE.set_progress_bar_config(disable=False) print("Wan VACE 14B BF16 loaded.") if SAM2_PREDICTOR is None: if build_sam2_video_predictor is None: raise RuntimeError( "SAM 2 is not installed. Install the official package with " "pip install git+https://github.com/facebookresearch/sam2.git. " f"Original import error: {_SAM2_IMPORT_ERROR!r}" ) print("Loading SAM 2.1 video predictor...") SAM2_PREDICTOR = build_sam2_video_predictor( SAM2_CONFIG, str(SAM2_CHECKPOINT), device=DEVICE ) print("SAM 2.1 loaded.") def make_difference_mask(original_frame, edited_reference, threshold=22): """Estimate the edited region from aligned original and AI-edited first frames.""" original = np.asarray(original_frame.convert("RGB"), dtype=np.int16) edited = np.asarray(edited_reference.convert("RGB").resize(original_frame.size, Image.Resampling.LANCZOS), dtype=np.int16) diff = np.max(np.abs(original - edited), axis=2).astype(np.uint8) # Threshold modest image-model noise, then clean isolated pixels and fill small holes. mask = (diff >= int(threshold)).astype(np.uint8) * 255 mask = cv2.medianBlur(mask, 5) kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9)) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2) mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)), iterations=1) ratio = float(np.count_nonzero(mask)) / float(mask.size) if ratio < 0.001: raise RuntimeError("The edited reference differs too little from the video's first frame to infer a reliable edit mask. Make sure it is the edited version of that exact frame.") if ratio > 0.80: raise RuntimeError("The inferred edit mask covers over 80% of the frame. Check that the reference has the same framing and that only the intended region was edited.") return (mask > 0).astype(np.uint8) def track_mask_from_reference(video_path, seed_mask, expected_frames): """Seed SAM 2.1 with the edited-region mask and propagate it through the video.""" if SAM2_PREDICTOR is None: raise RuntimeError("SAM 2.1 predictor is not loaded.") frame_dir = Path(video_path).with_suffix("") frame_dir.mkdir(parents=True, exist_ok=True) cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): raise RuntimeError(f"Could not open segmentation video: {video_path}") idx = 0 while True: ok, frame = cap.read() if not ok: break if idx >= int(expected_frames): break cv2.imwrite(str(frame_dir / f"{idx:06d}.jpg"), frame) idx += 1 cap.release() if idx == 0: raise RuntimeError("No frames found for SAM 2.1 tracking.") inference_state = SAM2_PREDICTOR.init_state(video_path=str(frame_dir)) try: SAM2_PREDICTOR.reset_state(inference_state) seed = cv2.resize(seed_mask.astype(np.uint8), (seed_mask.shape[1], seed_mask.shape[0]), interpolation=cv2.INTER_NEAREST) _, object_ids, mask_logits = SAM2_PREDICTOR.add_new_mask( inference_state=inference_state, frame_idx=0, obj_id=1, mask=seed.astype(bool), ) first_mask = (mask_logits[0, 0] > 0.0).detach().cpu().numpy().astype(np.uint8) tracked = {0: first_mask} for frame_idx, out_obj_ids, out_mask_logits in SAM2_PREDICTOR.propagate_in_video(inference_state): for j, obj_id in enumerate(out_obj_ids): if int(obj_id) == 1: tracked[int(frame_idx)] = (out_mask_logits[j, 0] > 0.0).detach().cpu().numpy().astype(np.uint8) break height, width = first_mask.shape output = [] previous = first_mask for i in range(int(expected_frames)): mask = tracked.get(i, previous) if mask.shape != (height, width): mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST) previous = mask output.append(mask) return output finally: try: SAM2_PREDICTOR.reset_state(inference_state) except Exception as e: print("SAM 2 state cleanup warning:", repr(e)) shutil.rmtree(frame_dir, ignore_errors=True) # ============================================================ # FACE PROTECTION # ============================================================ FACE_CASCADE = cv2.CascadeClassifier( cv2.data.haarcascades + "haarcascade_frontalface_default.xml" ) def face_mask(frame): rgb = np.asarray( frame.convert("RGB") ) gray = cv2.cvtColor( rgb, cv2.COLOR_RGB2GRAY, ) faces = FACE_CASCADE.detectMultiScale( gray, scaleFactor=1.1, minNeighbors=5, minSize=(40, 40), ) mask = np.zeros( gray.shape, dtype=np.uint8, ) for ( x, y, w, h, ) in faces: pad_x = int( w * 0.20 ) pad_y = int( h * 0.25 ) x1 = max( 0, x - pad_x, ) y1 = max( 0, y - pad_y, ) x2 = min( gray.shape[1], x + w + pad_x, ) y2 = min( gray.shape[0], y + h + pad_y, ) mask[ y1:y2, x1:x2, ] = 1 # Additional dilation. kernel = np.ones( (11, 11), dtype=np.uint8, ) mask = cv2.dilate( mask, kernel, iterations=1, ) return mask # ============================================================ # FINAL EDIT MASK # ============================================================ def create_edit_masks( target_masks, frames, ): output = [] for target, frame in zip( target_masks, frames, ): protected = face_mask( frame ) target = target.astype( np.uint8 ) # Target minus protected face. editable = target * ( 1 - protected ) # Feather boundary. editable = cv2.GaussianBlur( editable.astype( np.float32 ), (0, 0), sigmaX=3.0, ) editable = np.clip( editable, 0.0, 1.0, ) output.append( editable ) return output # ============================================================ # PIXEL COMPOSITING # ============================================================ def _frame_to_rgb_array(frame): """ Normalize PIL / NumPy / torch outputs into HxWx3 uint8. Wan VACE may return NumPy arrays, while other versions can return PIL. """ if isinstance(frame, Image.Image): arr = np.asarray(frame.convert("RGB")) elif torch.is_tensor(frame): arr = frame.detach().float().cpu().numpy() else: arr = np.asarray(frame) if arr.ndim == 4 and arr.shape[0] == 1: arr = arr[0] # Handle CHW tensors/arrays. if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): arr = np.transpose(arr, (1, 2, 0)) if arr.ndim == 2: arr = np.stack([arr, arr, arr], axis=-1) if arr.shape[-1] == 4: arr = arr[..., :3] if np.issubdtype(arr.dtype, np.floating): finite_max = float(np.nanmax(arr)) if arr.size else 0.0 if finite_max <= 1.5: arr = arr * 255.0 return np.clip(arr, 0, 255).astype(np.uint8) def composite( original, generated, mask, ): original = _frame_to_rgb_array(original).astype(np.float32) generated = _frame_to_rgb_array(generated).astype(np.float32) if generated.shape != original.shape: generated = cv2.resize( generated, ( original.shape[1], original.shape[0], ), interpolation=cv2.INTER_LANCZOS4, ) mask = cv2.resize( mask.astype(np.float32), ( original.shape[1], original.shape[0], ), interpolation=cv2.INTER_LINEAR, ) mask = np.clip(mask, 0.0, 1.0)[..., None] result = ( generated * mask + original * (1.0 - mask) ) return Image.fromarray( np.clip( result, 0, 255, ).astype(np.uint8) ) def mask_preview(mask): """Convert a 0..1 or 0/1 edit mask to a visible grayscale PIL image.""" arr = np.asarray(mask, dtype=np.float32) if arr.ndim == 3: arr = arr[..., 0] arr = np.clip(arr, 0.0, 1.0) return Image.fromarray((arr * 255).astype(np.uint8), mode="L").convert("RGB") def masked_difference_stats(original, generated, mask): """Report whether generation materially changed the masked area.""" a = _frame_to_rgb_array(original).astype(np.float32) b = _frame_to_rgb_array(generated).astype(np.float32) if b.shape != a.shape: b = cv2.resize(b, (a.shape[1], a.shape[0]), interpolation=cv2.INTER_LANCZOS4) m = np.asarray(mask, dtype=np.float32) if m.shape != a.shape[:2]: m = cv2.resize(m, (a.shape[1], a.shape[0]), interpolation=cv2.INTER_LINEAR) selected = m > 0.25 outside = ~selected if not np.any(selected): return {"mask_pct": 0.0, "inside_mae": 0.0, "outside_mae": float(np.mean(np.abs(a-b)))} inside_mae = float(np.mean(np.abs(a[selected] - b[selected]))) outside_mae = float(np.mean(np.abs(a[outside] - b[outside]))) if np.any(outside) else 0.0 return {"mask_pct": float(np.mean(selected) * 100.0), "inside_mae": inside_mae, "outside_mae": outside_mae} # ============================================================ # WAN GENERATION # ============================================================ def generate_wan( frames, prompt, edit_masks, reference, strength, steps, guidance, seed, ): generator = torch.Generator( device=DEVICE ).manual_seed( int(seed) ) # Wan VACE is the native Wan 2.1 pipeline for masked # video-to-video editing plus optional reference-image # conditioning. Black mask pixels are kept as conditioning; # white pixels are generated/edited. vace_masks = [ Image.fromarray( np.clip( mask * 255.0, 0, 255, ).astype(np.uint8), mode="L", ) for mask in edit_masks ] kwargs = dict( video=frames, mask=vace_masks, prompt=prompt, negative_prompt=( "low quality, worst quality, blurry, " "distorted face, deformed face, " "bad anatomy, extra limbs, " "warped hands, flickering, " "frame artifacts, subtitles, " "watermark, text, nude, naked, topless, " "bare breasts, exposed chest, exposed groin" ), height=frames[0].height, width=frames[0].width, num_frames=len(frames), num_inference_steps=int( steps ), guidance_scale=float( guidance ), conditioning_scale=float(strength), generator=generator, ) if reference is not None: kwargs["reference_images"] = [ reference ] # In WanVACEPipeline, conditioning_scale weights the VACE control # stream. It is NOT an img2img denoising strength or a pixel blend. # The final composite uses the generated result throughout the edit mask. with torch.inference_mode(): result = WAN_PIPE( **kwargs ) return result.frames[0] # ============================================================ # VACE CHUNKING # ============================================================ def _vace_model_length(frame_count): """Round a requested chunk up to Wan's 4n+1 temporal shape.""" frame_count = max(5, int(frame_count)) return ((frame_count - 1 + 3) // 4) * 4 + 1 def build_vace_chunks(frames, masks): """Yield overlapping source chunks without dropping video duration.""" total = len(frames) step = VACE_CHUNK_FRAMES - VACE_CHUNK_OVERLAP if total <= VACE_CHUNK_FRAMES: yield 0, frames, masks, total return start = 0 while start < total: end = min(start + VACE_CHUNK_FRAMES, total) desired = end - start model_len = _vace_model_length(desired) chunk_frames = list(frames[start:end]) chunk_masks = list(masks[start:end]) # Pad only the short final chunk. These padding frames are # discarded again after generation, so the output frame count # exactly matches the source working-rate frame count. while len(chunk_frames) < model_len: chunk_frames.append(chunk_frames[-1].copy()) chunk_masks.append(chunk_masks[-1].copy()) yield start, chunk_frames, chunk_masks, desired if end >= total: break start += step # ============================================================ # ENCODE # ============================================================ def encode_video( frames, fps, output, ): # Diffusers export helper. export_to_video( frames, str(output), fps=fps, ) return str(output) # ============================================================ # AUDIO REMUX # ============================================================ def remux_audio( silent_video, original_video, output, ): # We take the ORIGINAL audio directly. # # Video is re-encoded only once here. # Audio remains untouched. run_command([ "ffmpeg", "-y", "-i", silent_video, "-i", original_video, "-map", "0:v:0", "-map", "1:a:0?", "-c:v", "libx264", "-preset", "medium", "-crf", "16", "-pix_fmt", "yuv420p", "-c:a", "copy", "-shortest", output, ]) return str(output) # ============================================================ # MAIN GPU FUNCTION # ============================================================ @spaces.GPU( size="xlarge", # Override with GPU_DURATION if your Space quota permits a longer job. duration=int(os.environ.get("GPU_DURATION", "360")), ) def process_video( video, reference, prompt, strength, steps, guidance, seed, ): video_path = get_video_path( video ) if video_path is None: raise gr.Error( "Upload a video." ) if not prompt: raise gr.Error( "Enter an editing prompt." ) try: # ---------------------------------------------------- # MODEL INITIALIZATION # ---------------------------------------------------- load_models() # ---------------------------------------------------- # INPUT VALIDATION # ---------------------------------------------------- info = probe_video( video_path ) if ( info["duration"] > MAX_SECONDS ): raise gr.Error( f"Maximum video length is " f"{MAX_SECONDS} seconds." ) # ---------------------------------------------------- # WORKSPACE # ---------------------------------------------------- job_id = uuid.uuid4().hex work = ( OUTPUT_DIR / job_id ) work.mkdir( parents=True, exist_ok=True, ) # ---------------------------------------------------- # AUDIO # ---------------------------------------------------- audio_path = ( work / "audio.mka" ) if info["has_audio"]: extract_audio( video_path, str(audio_path), ) # ---------------------------------------------------- # FRAMES # ---------------------------------------------------- frames, fps, _ = decode_frames( video_path ) frames = resize_for_wan( frames ) # ---------------------------------------------------- # SOURCE FRAME COUNT # ---------------------------------------------------- # Do not truncate a longer source to 81 frames. VACE is # processed in overlapping <=81-frame chunks so the final # video retains the full source duration. if len(frames) < 5: raise gr.Error( "Video is too short for Wan VACE." ) # ---------------------------------------------------- # EDITED FIRST-FRAME DIFFERENCE + SAM 2.1 TRACKING # ---------------------------------------------------- if reference is None: raise gr.Error("Upload the AI-edited version of the video's first frame as the reference image.") reference = reference.convert("RGB").resize(frames[0].size, Image.Resampling.LANCZOS) print("Inferring edit mask by comparing original and edited first frames...") seed_mask = make_difference_mask(frames[0], reference) sam_video = work / "sam_working.mp4" encode_video(frames, MODEL_FPS, sam_video) print("Propagating first-frame edit mask with SAM 2.1...") tracked_masks = track_mask_from_reference(sam_video, seed_mask, len(frames)) # ---------------------------------------------------- # FACE PROTECTION # ---------------------------------------------------- print( "Building protected regions..." ) edit_masks = ( create_edit_masks( tracked_masks, frames, ) ) # ---------------------------------------------------- # DIAGNOSTIC ARTIFACTS: inspect mask before paying for generation # ---------------------------------------------------- diagnostics_dir = work / "diagnostics" diagnostics_dir.mkdir(parents=True, exist_ok=True) frames[0].save(diagnostics_dir / "01_original_first_frame.png") reference.save(diagnostics_dir / "02_edited_reference.png") mask_preview(seed_mask).save(diagnostics_dir / "03_raw_difference_mask.png") mask_preview(tracked_masks[0]).save(diagnostics_dir / "04_sam2_tracked_mask_first_frame.png") mask_preview(edit_masks[0]).save(diagnostics_dir / "05_final_vace_mask_first_frame.png") seed_pct = float(np.mean(seed_mask > 0) * 100.0) tracked_pct = float(np.mean(np.asarray(tracked_masks[0]) > 0) * 100.0) final_pct = float(np.mean(np.asarray(edit_masks[0]) > 0.25) * 100.0) print(f"MASK DIAGNOSTICS: raw_diff={seed_pct:.2f}% tracked={tracked_pct:.2f}% final_edit={final_pct:.2f}%") if final_pct < 0.5: raise RuntimeError( f"Final edit mask covers only {final_pct:.2f}% of the frame. " "Inspect diagnostics/03-05 mask previews; the selected region is too small to edit reliably." ) # ---------------------------------------------------- # PROMPT # ---------------------------------------------------- final_prompt = prompt if reference is not None: final_prompt += ( ". CLOTHING EDIT: the adult subject must be wearing " "an opaque two-piece bikini matching the supplied edited " "reference image. The bikini top must cover the breasts and " "the bikini bottom must cover the groin. Replace the original " "nudity with the bikini; do not leave the original nude " "appearance unchanged. Preserve the subject's identity, face, " "hair, pose, hands, camera framing, lighting, and background " "as closely as possible. Treat the edited reference as the " "target appearance for the clothing." ) # ---------------------------------------------------- # WAN VACE — FULL VIDEO, CHUNKED # ---------------------------------------------------- print( f"Running Wan VACE 14B in <= {VACE_CHUNK_FRAMES}-frame chunks..." ) final_frames = [] first_chunk = True first_generated_preview = None first_composite_preview = None first_chunk_stats = None for chunk_index, ( start, chunk_frames, chunk_masks, desired_count, ) in enumerate( build_vace_chunks(frames, edit_masks) ): print( f"VACE chunk {chunk_index + 1}: " f"source frames {start}..{start + desired_count - 1} " f"(model length={len(chunk_frames)})" ) generated = generate_wan( frames=chunk_frames, prompt=final_prompt, edit_masks=chunk_masks, reference=reference, strength=float(strength), steps=int(steps), guidance=float(guidance), seed=int(seed), ) generated = list(generated[:desired_count]) if chunk_index == 0 and generated: first_generated_preview = Image.fromarray(_frame_to_rgb_array(generated[0])) first_generated_preview.save(diagnostics_dir / "06_vace_generated_first_frame.png") first_chunk_stats = masked_difference_stats(frames[0], generated[0], chunk_masks[0]) print( "GENERATION DIAGNOSTICS: " f"mask={first_chunk_stats['mask_pct']:.2f}% " f"mean_abs_diff_inside={first_chunk_stats['inside_mae']:.2f}/255 " f"outside={first_chunk_stats['outside_mae']:.2f}/255" ) if first_chunk_stats["inside_mae"] < 3.0: print("WARNING: VACE changed the masked region very little. Check prompt, mask preview, and reference conditioning.") # Preserve the exact source frame count. Overlap frames are # used as temporal context for VACE but emitted only once. if first_chunk: emit_from = 0 first_chunk = False else: emit_from = min(VACE_CHUNK_OVERLAP, len(generated)) for local_idx in range(emit_from, len(generated)): global_idx = start + local_idx if global_idx >= len(frames): break composited = composite( frames[global_idx], generated[local_idx], np.clip(chunk_masks[local_idx], 0.0, 1.0), ) if global_idx == 0: first_composite_preview = composited composited.save(diagnostics_dir / "07_final_composite_first_frame.png") final_frames.append(composited) if len(final_frames) != len(frames): raise RuntimeError( f"Internal frame accounting error: " f"expected {len(frames)} frames, got {len(final_frames)}." ) # ---------------------------------------------------- # DIAGNOSTIC REPORT + DOWNLOADABLE ZIP # ---------------------------------------------------- report = { "source_video": str(video_path), "frame_count": len(frames), "working_fps": float(fps), "working_resolution": [frames[0].width, frames[0].height], "raw_difference_mask_percent": round(seed_pct, 3), "sam2_first_frame_mask_percent": round(tracked_pct, 3), "final_edit_mask_percent": round(final_pct, 3), "vace_conditioning_scale": float(strength), "inference_steps": int(steps), "guidance_scale": float(guidance), "user_prompt": str(prompt), "effective_prompt": str(final_prompt), "negative_prompt": ( "low quality, worst quality, blurry, distorted face, deformed face, " "bad anatomy, extra limbs, warped hands, flickering, frame artifacts, " "subtitles, watermark, text, nude, naked, topless, bare breasts, " "exposed chest, exposed groin" ), "first_frame_generation_stats": first_chunk_stats, "notes": [ "Wan VACE mask convention: white means generate/edit; black means preserve/condition.", "VACE conditioning_scale controls the VACE conditioning stream; it is not an img2img strength parameter.", "The generated preview is saved before final masked compositing.", ], } (diagnostics_dir / "report.json").write_text(json.dumps(report, indent=2)) diagnostics_zip = work / "diagnostics.zip" with zipfile.ZipFile(diagnostics_zip, "w", compression=zipfile.ZIP_DEFLATED) as zf: for artifact in sorted(diagnostics_dir.iterdir()): if artifact.is_file(): zf.write(artifact, arcname=artifact.name) # ---------------------------------------------------- # VIDEO # ---------------------------------------------------- silent_video = ( work / "video.mp4" ) encode_video( final_frames, fps, silent_video, ) # ---------------------------------------------------- # AUDIO # ---------------------------------------------------- final_video = ( work / "final.mp4" ) if info["has_audio"]: remux_audio( str(silent_video), video_path, str(final_video), ) else: shutil.copy2( silent_video, final_video, ) print( "Done:", final_video, ) # Return video plus previews and a zip of diagnostic artifacts. seed_preview = mask_preview(seed_mask) sam_preview = mask_preview(tracked_masks[0]) final_mask_preview = mask_preview(edit_masks[0]) generated_preview = first_generated_preview or frames[0] composite_preview = first_composite_preview or frames[0] return ( str(final_video), seed_preview, sam_preview, final_mask_preview, generated_preview, composite_preview, str(diagnostics_zip), ) except gr.Error: raise except Exception as e: print( "PIPELINE ERROR:", repr(e), ) raise gr.Error( f"{type(e).__name__}: {e}" ) finally: cleanup() # ============================================================ # DOWNLOAD MODELS BEFORE THE APP STARTS # ============================================================ # This is deliberately outside @spaces.GPU. # Network/model download does not need a GPU and happens once at startup. ensure_model_downloads() # ============================================================ # UI # ============================================================ CSS = """ #title { text-align: center; font-size: 34px; font-weight: 700; } #subtitle { text-align: center; opacity: 0.7; } .gradio-container { max-width: 1400px !important; } """ with gr.Blocks( title="Zero-Drift VACE Auto Editor" ) as demo: gr.Markdown( "# Zero-Drift VACE Auto Editor", elem_id="title", ) gr.Markdown( "SAM 2.1 mask propagation + Wan 2.1 VACE 14B BF16 + full-duration chunked editing + diagnostics", elem_id="subtitle", ) with gr.Row(): # ==================================================== # LEFT # ==================================================== with gr.Column( scale=1 ): video_input = gr.Video( label="Source Video" ) reference_input = gr.Image( label="Edited First Frame (required)", type="pil", ) prompt_input = gr.Textbox( label="Edit Instruction", placeholder=( "Describe the intended clothing edit shown in the reference. " "Example: Replace nudity with an opaque two-piece bikini; preserve identity, pose, lighting, and background." ), lines=4, ) strength_input = gr.Slider( minimum=0.50, maximum=1.50, value=1.00, step=0.05, label="VACE Conditioning Scale", info="Weights VACE's control-conditioning stream. This is not a pixel-blend slider; inspect the raw generated preview and mask diagnostics.", ) steps_input = gr.Slider( minimum=8, maximum=13, value=12, step=1, label="Inference Steps (higher values take longer)", info="Default raised from 10 to 12 because the diagnostic run barely changed the masked area. If the job hits its runtime limit, use a shorter clip.", ) guidance_input = gr.Slider( minimum=1.0, maximum=8.0, value=6.0, step=0.5, label="Guidance", ) seed_input = gr.Number( value=42, precision=0, label="Seed", ) generate_button = gr.Button( "GENERATE", variant="primary", ) # ==================================================== # RIGHT # ==================================================== with gr.Column( scale=1 ): gr.Markdown( "### Edited-first-frame workflow\n\n" "Upload the original video and an edited image of its exact first frame. The app estimates the changed region by comparing the two images, then SAM 2.1 tracks that region through the video.\n\n" "The prompt should explicitly describe the intended clothing result. Review the raw mask and raw VACE output previews: if the mask is wrong, fix the reference/mask; if the mask is right but the raw VACE frame is weak, generation conditioning is the issue. The reference must match the original first frame's framing." ) result_video = gr.Video( label="Result" ) gr.Markdown("### Diagnostics (use these to identify the failing stage)") mask_preview_output = gr.Image(label="Raw difference mask (white = changed)") sam_mask_preview_output = gr.Image(label="SAM 2.1 propagated mask") final_mask_preview_output = gr.Image(label="Final VACE edit mask (white = regenerate)") generated_preview_output = gr.Image(label="Raw VACE first-frame output, before compositing") composite_preview_output = gr.Image(label="Final first-frame composite") diagnostics_file_output = gr.File(label="Download diagnostics ZIP") # ======================================================== # GENERATE # ======================================================== generate_button.click( fn=process_video, inputs=[ video_input, reference_input, prompt_input, strength_input, steps_input, guidance_input, seed_input, ], outputs=[ result_video, mask_preview_output, sam_mask_preview_output, final_mask_preview_output, generated_preview_output, composite_preview_output, diagnostics_file_output, ], concurrency_limit=1, ) # ============================================================ # LAUNCH # ============================================================ if __name__ == "__main__": # Download model artifacts BEFORE Gradio starts. # This means a user request can never trigger a fresh Wan download, # accepts work. ensure_model_downloads() print("All startup model downloads are complete.") print("Starting Gradio...") # IMPORTANT: HF Spaces may enable Gradio SSR by default. # For this long-running GPU job, SSR/Node proxying can misroute # heartbeat traffic as component payloads (e.g. "ping" -> FileData). # Disable SSR and use the normal client-rendered Gradio frontend. demo.queue( max_size=1, ).launch( css=CSS, ssr_mode=False, )