copy of thornmaze/wan2-2-i2v-v3 (byte-1:1)

README.md

@@ -1,11 +1,10 @@
 ---
-title: Wan2 2 I2v V3
-emoji: 🐢
-colorFrom: blue
-colorTo: purple
+title: Wan2.2 14B Fast Preview [NEW]
+emoji: 🏆
+colorFrom: indigo
+colorTo: pink
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 6.0.1
 app_file: app.py
 pinned: false
 ---

aoti.py

@@ -0,0 +1,35 @@
+"""
+"""
+
+from typing import cast
+
+import torch
+from huggingface_hub import hf_hub_download
+from spaces.zero.torch.aoti import ZeroGPUCompiledModel
+from spaces.zero.torch.aoti import ZeroGPUWeights
+from torch._functorch._aot_autograd.subclass_parametrization import unwrap_tensor_subclass_parameters
+
+
+def _shallow_clone_module(module: torch.nn.Module) -> torch.nn.Module:
+    clone = object.__new__(module.__class__)
+    clone.__dict__ = module.__dict__.copy()
+    clone._parameters = module._parameters.copy()
+    clone._buffers = module._buffers.copy()
+    clone._modules = {k: _shallow_clone_module(v) for k, v in module._modules.items() if v is not None}
+    return clone
+
+
+def aoti_blocks_load(module: torch.nn.Module, repo_id: str, variant: str | None = None):
+    repeated_blocks = cast(list[str], module._repeated_blocks)
+    aoti_files = {name: hf_hub_download(
+        repo_id=repo_id,
+        filename='package.pt2',
+        subfolder=name if variant is None else f'{name}.{variant}',
+    ) for name in repeated_blocks}
+    for block_name, aoti_file in aoti_files.items():
+        for block in module.modules():
+            if block.__class__.__name__ == block_name:
+                block_ = _shallow_clone_module(block)
+                unwrap_tensor_subclass_parameters(block_)
+                weights = ZeroGPUWeights(block_.state_dict())
+                block.forward = ZeroGPUCompiledModel(aoti_file, weights)

app.py

@@ -0,0 +1,725 @@
+import os; os.system('pip install --upgrade --no-deps spaces')
+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
+import torch._dynamo
+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
+import lora_loader
+
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+warnings.filterwarnings("ignore")
+IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))
+
+# if IS_ZERO_GPU:
+#     print("Loading...")
+#     subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+
+# --- FRAME EXTRACTION JS & LOGIC ---
+
+# JS to grab timestamp from the output video
+get_timestamp_js = """
+function() {
+    // Select the video element specifically inside the component with id 'generated-video'
+    const video = document.querySelector('#generated-video video');
+    
+    if (video) {
+        console.log("Video found! Time: " + video.currentTime);
+        return video.currentTime;
+    } else {
+        console.log("No video element found.");
+        return 0;
+    }
+}
+"""
+
+
+def extract_frame(video_path, timestamp):
+    # Safety check: if no video is present
+    if not video_path:
+        return None
+    
+    print(f"Extracting frame at timestamp: {timestamp}") 
+    
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        return None
+
+    # Calculate frame number
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    target_frame_num = int(float(timestamp) * fps)
+    
+    # Cap total frames to prevent errors at the very end of video
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if target_frame_num >= total_frames:
+        target_frame_num = total_frames - 1
+    
+    # Set position
+    cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame_num)
+    ret, frame = cap.read()
+    cap.release()
+    
+    if ret:
+        # Convert from BGR (OpenCV) to RGB (Gradio)
+        # Gradio Image component handles Numpy array -> PIL conversion automatically
+        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    
+    return None
+
+# --- END FRAME EXTRACTION LOGIC ---
+
+
+def clear_vram():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+# RIFE
+if not os.path.exists("RIFEv4.26_0921.zip"):
+    print("Downloading RIFE Model...")
+    subprocess.run([
+        "wget", "-q",
+        "https://huggingface.co/thornmaze/RIFE/resolve/main/RIFEv4.26_0921.zip",
+        "-O", "RIFEv4.26_0921.zip"
+    ], check=True)
+    subprocess.run(["unzip", "-o", "RIFEv4.26_0921.zip"], check=True)
+
+# sys.path.append(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()
+
+
+@torch.no_grad()
+def interpolate_bits(frames_np, multiplier=2, scale=1.0):
+    """
+    Interpolation maintaining Numpy Float 0-1 format.
+    Args:
+        frames_np: Numpy Array (Time, Height, Width, Channels) - Float32 [0.0, 1.0]
+        multiplier: int (2, 4, 8)
+    Returns:
+        List of Numpy Arrays (Height, Width, Channels) - Float32 [0.0, 1.0]
+    """
+    
+    # Handle input shape
+    if isinstance(frames_np, list):
+        # Convert list of arrays to one big array for easier shape handling if needed, 
+        # but here we just grab dims from first frame
+        T = len(frames_np)
+        H, W, C = frames_np[0].shape
+    else:
+        T, H, W, C = frames_np.shape
+
+    # 1. No Interpolation Case
+    if multiplier < 2:
+        # Just convert 4D array to list of 3D arrays
+        if isinstance(frames_np, np.ndarray):
+            return list(frames_np)
+        return frames_np
+
+    n_interp = multiplier - 1
+    
+    # Pre-calc padding for RIFE (requires dimensions divisible by 32/scale)
+    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)
+
+    # Helper: Numpy (H, W, C) Float -> Tensor (1, C, H, W) Half
+    def to_tensor(frame_np):
+        # frame_np is float32 0-1
+        t = torch.from_numpy(frame_np).to(device)
+        # HWC -> CHW
+        t = t.permute(2, 0, 1).unsqueeze(0)
+        return F.pad(t, padding).half()
+
+    # Helper: Tensor (1, C, H, W) Half -> Numpy (H, W, C) Float
+    def from_tensor(tensor):
+        # Crop padding
+        t = tensor[0, :, :H, :W]
+        # CHW -> HWC
+        t = t.permute(1, 2, 0)
+        # Keep as float32, range 0-1
+        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 = []
+
+    # Process Frames
+    # Load first frame into GPU
+    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
+            # Add original frame to output
+            output_frames.append(from_tensor(I0))
+    
+            # Load next frame
+            I1 = to_tensor(frames_np[i+1])
+    
+            # Generate intermediate frames
+            mid_tensors = make_inference(I0, I1, n_interp)
+    
+            # Append intermediate frames
+            for mid in mid_tensors:
+                output_frames.append(from_tensor(mid))
+
+            if (i + 1) % 50 == 0:
+                pbar.update(50)
+        pbar.update(total_steps % 50)
+        
+        # Add the very last frame
+        output_frames.append(from_tensor(I1))
+    
+    # Cleanup
+    del I0, I1, mid_tensors
+    torch.cuda.empty_cache()
+
+    return output_frames
+
+
+# WAN
+
+MODEL_ID = "thornmaze/WAMU_v3_WAN2.2_I2V_LIGHTNING"
+
+LORA_MODELS = []
+
+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 = 321
+
+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,
+}
+
+pipe = WanImageToVideoPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+).to('cuda')
+original_scheduler = copy.deepcopy(pipe.scheduler)
+
+for i, lora in enumerate(LORA_MODELS):
+    name_high_tr = lora["high_tr"].split(".")[0].split("/")[-1] + "Hh"
+    name_low_tr = lora["low_tr"].split(".")[0].split("/")[-1] + "Ll"
+    
+    try: 
+        pipe.load_lora_weights(
+            lora["repo_id"],
+            weight_name=lora["high_tr"],
+            adapter_name=name_high_tr
+        )
+    
+        kwargs_lora = {"load_into_transformer_2": True}
+        pipe.load_lora_weights(
+            lora["repo_id"],
+            weight_name=lora["low_tr"],
+            adapter_name=name_low_tr,
+            **kwargs_lora
+        )
+    
+        pipe.set_adapters([name_high_tr, name_low_tr], adapter_weights=[1.0, 1.0])
+    
+        pipe.fuse_lora(adapter_names=[name_high_tr], lora_scale=lora["high_scale"], components=["transformer"])
+        pipe.fuse_lora(adapter_names=[name_low_tr], lora_scale=lora["low_scale"], components=["transformer_2"])
+    
+        pipe.unload_lora_weights()
+
+        print(f"Applied: {lora['high_tr']}, hs={lora['high_scale']}/ls={lora['low_scale']}, {i+1}/{len(LORA_MODELS)}") 
+    except Exception as e:
+        print("Error:", str(e))
+        print("Failed LoRA:", name_high_tr)
+        pipe.unload_lora_weights()
+
+# if os.path.exists(CACHE_DIR):
+#     shutil.rmtree(CACHE_DIR)
+#     print("Deleted Hugging Face cache.")
+# else:
+#     print("No hub cache found.")
+
+quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
+torch._dynamo.reset()
+quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+torch._dynamo.reset()
+quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
+torch._dynamo.reset()
+
+spaces.aoti_load(
+    module=pipe.transformer,
+    repo_id='thornmaze/WanTransformer3DModel-sm120-cu130-raa',
+)
+spaces.aoti_load(
+    module=pipe.transformer_2,
+    repo_id='thornmaze/WanTransformer3DModel-sm120-cu130-raa',
+)
+
+# pipe.vae.enable_slicing()
+# pipe.vae.enable_tiling()
+
+default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
+default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
+
+
+def model_title():
+    return "## Wan 2.2 I2V 14B Lightning — NSFW"
+
+
+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):
+    raw = int(round(duration_seconds * FIXED_FPS))
+    raw = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, raw))
+    return ((raw - 1) // 4) * 4 + 1
+
+
+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,
+    safe_mode,
+    lora_groups,
+    progress
+):
+    BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624
+    BASE_STEP_DURATION = 5.
+    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 * 2.4
+
+    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
+
+    total_time = 15 + gen_time
+    if safe_mode:
+        total_time = total_time * 1.30
+
+    return total_time
+
+
+@spaces.GPU(duration=get_inference_duration, size='xlarge')
+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,
+    safe_mode=False,
+    lora_groups=None,
+    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"Generating {num_frames} frames, task: {task_name}, {duration_seconds}, {resized_image.size}, lora={lora_groups}")
+    start = time.time()
+
+    lora_loaded = False
+    if lora_groups:
+        try:
+            for idx, name in enumerate(lora_groups):
+                if name and name != "(None)":
+                    lora_loader.load_lora_to_pipe(pipe, name, adapter_name=f"lora_{idx}")
+            lora_loaded = True
+            print(f"LoRA loaded: {lora_groups}")
+        except Exception as e:
+            print(f"LoRA warning: {e}")
+
+    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"
+    )
+
+    if lora_loaded:
+        lora_loader.unload_lora(pipe)
+
+    print("gen time passed:", time.time() - start)
+
+    raw_frames_np = result.frames[0]  # Returns (T, H, W, C) float32
+    pipe.scheduler = original_scheduler
+
+    frame_factor = frame_multiplier // FIXED_FPS
+    if frame_factor > 1:
+        start = time.time()
+        print(f"Processing frames (RIFE Multiplier: {frame_factor}x)...")
+        rife_model.device()
+        rife_model.flownet = rife_model.flownet.half()
+        final_frames = interpolate_bits(raw_frames_np, multiplier=int(frame_factor))
+        print("Interpolation time passed:", time.time() - start)
+    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
+
+    start = time.time()
+    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)
+    print(f"Export time passed, {final_fps} FPS:", time.time() - start)
+
+    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,
+    safe_mode=False,
+    lora_groups=None,
+    video_component=True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    Generate a video from an input image using the Wan 2.2 14B I2V model with Lightning LoRA.
+    This function takes an input image and generates a video animation based on the provided
+    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Image-to-Video model in with Lightning LoRA
+    for fast generation in 4-8 steps.
+    Args:
+        input_image (PIL.Image): The input image to animate. Will be resized to target dimensions.
+        last_image (PIL.Image, optional): The optional last image for the video.
+        prompt (str): Text prompt describing the desired animation or motion.
+        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
+            Defaults to 4. Range: 1-30.
+        negative_prompt (str, optional): Negative prompt to avoid unwanted elements.
+            Defaults to default_negative_prompt (contains unwanted visual artifacts).
+        duration_seconds (float, optional): Duration of the generated video in seconds.
+            Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
+        guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
+            Defaults to 1.0. Range: 0.0-20.0.
+        guidance_scale_2 (float, optional): Controls adherence to the prompt. Higher values = more adherence.
+            Defaults to 1.0. Range: 0.0-20.0.
+        seed (int, optional): Random seed for reproducible results. Defaults to 42.
+            Range: 0 to MAX_SEED (2147483647).
+        randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
+            Defaults to False.
+        quality (float, optional): Video output quality. Default is 5. Uses variable bit rate.
+            Highest quality is 10, lowest is 1.
+        scheduler (str, optional): The name of the scheduler to use for inference. Defaults to "UniPCMultistep".
+        flow_shift (float, optional): The flow shift value for compatible schedulers. Defaults to 6.0.
+        frame_multiplier (int, optional): The int value for fps enhancer
+        video_component(bool, optional): Show video player in output.
+            Defaults to True.
+        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
+    Returns:
+        tuple: A tuple containing:
+            - video_path (str): Path for the video component.
+            - video_path (str): Path for the file download component. Attempt to avoid reconversion in video component.
+            - current_seed (int): The seed used for generation.
+    Raises:
+        gr.Error: If input_image is None (no image uploaded).
+    Note:
+        - Frame count is calculated as duration_seconds * FIXED_FPS (24)
+        - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
+        - The function uses GPU acceleration via the @spaces.GPU decorator
+        - Generation time varies based on steps and duration (see get_duration function)
+    """
+    
+    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,
+        safe_mode,
+        lora_groups,
+        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(delete_cache=(3600, 10800)) as demo:
+    gr.Markdown(model_title())
+    gr.Markdown("Run Wan 2.2 in just 4-8 steps, fp8 quantization & AoT compilation - 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, FIXED_FPS*8],
+                value=FIXED_FPS,
+                label="Video Fluidity (Frames per Second)",
+                info="Extra frames will be generated using flow estimation, which estimates motion between frames to make the video smoother."
+            )
+            safe_mode_checkbox = gr.Checkbox(
+                label="🛠️ Safe Mode",
+                value=True,
+                info="Requests 30% extra processing time to try to prevent unfinished tasks when the server is busy."
+            )
+            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", info="If set to 10, the generated video may be too large and won't play in the Gradio preview.")
+                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 - high noise stage", info="Values above 1 increase GPU usage and may take longer to process.")
+                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage")
+                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")
+                lora_dropdown = gr.Dropdown(choices=lora_loader.get_lora_choices(), label="LoRA (NSFW)", multiselect=True, info="Select scenario LoRAs")
+                play_result_video = gr.Checkbox(label="Display result", value=True, interactive=True)
+
+            generate_button = gr.Button("Generate Video", variant="primary")
+
+        with gr.Column():
+            # ASSIGNED elem_id="generated-video" so JS can find it
+            video_output = gr.Video(label="Generated Video", autoplay=True, sources=["upload"], buttons=["download", "share"], interactive=True, elem_id="generated-video")
+            
+            # --- Frame Grabbing UI ---
+            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,
+        safe_mode_checkbox,
+        lora_dropdown,
+        play_result_video
+    ]
+    
+    generate_button.click(
+        fn=generate_video, 
+        inputs=ui_inputs, 
+        outputs=[video_output, file_output, seed_input]
+    )
+    
+    # --- Frame Grabbing Events ---
+    # 1. Click button -> JS runs -> puts time in hidden number box
+    grab_frame_btn.click(
+        fn=None,
+        inputs=None,
+        outputs=[timestamp_box],
+        js=get_timestamp_js
+    )
+    
+    # 2. Hidden number box changes -> Python runs -> puts frame in Input Image
+    timestamp_box.change(
+        fn=extract_frame,
+        inputs=[video_output, timestamp_box],
+        outputs=[input_image_component]
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch(
+        mcp_server=True,
+        css=CSS,
+        show_error=True,
+    )

lora_loader.py

@@ -0,0 +1,201 @@
+"""
+LoRA Loader for WAN 2.2 I2V - multi-repo.
+Baseline: lkzd7/WAN2.2_LoraSet_NSFW. Plus curated I2V act/helper LoRAs from
+Lythiga/WAN2.2_I2V_Lora, rolechar/rc_wan2.2_i2v_loraset_nsfw_private,
+lopi999/Wan2.2-I2V_General-NSFW-LoRA. T2V / lightning / furry / dance / male-char filtered out.
+"""
+import os
+import urllib.parse
+import re
+from huggingface_hub import hf_hub_download
+
+LORA_REPO = "lkzd7/WAN2.2_LoraSet_NSFW"
+HF_TOKEN = os.environ.get("HF_TOKEN")  # authenticated downloads (covers private repos)
+
+# Pinned commit hashes (2026-07-19) — protection from breaking upstream changes
+PINNED_REVISIONS = {
+    "lkzd7/WAN2.2_LoraSet_NSFW": "ccfa867ce4b1",
+    "Lythiga/WAN2.2_I2V_Lora": "67ef542e6faa",
+    "rolechar/rc_wan2.2_i2v_loraset_nsfw_private": "4855a2f216ac",
+    "lopi999/Wan2.2-I2V_General-NSFW-LoRA": "aeef17d7fa51",
+}
+
+LORA_FILES = [
+    "Blink_Squatting_Cowgirl_Position_I2V_HIGH.safetensors",
+    "Blink_Squatting_Cowgirl_Position_I2V_LOW.safetensors",
+    "PENISLORA_22_i2v_HIGH_e320.safetensors",
+    "PENISLORA_22_i2v_LOW_e496.safetensors",
+    "Pornmaster_wan 2.2_14b_I2V_bukkake_v1.4_high_noise.safetensors",
+    "Pornmaster_wan 2.2_14b_I2V_bukkake_v1.4_low_noise.safetensors",
+    "W22_Multiscene_Photoshoot_Softcore_i2v_HN.safetensors",
+    "W22_Multiscene_Photoshoot_Softcore_i2v_LN.safetensors",
+    "WAN-2.2-I2V-Double-Blowjob-HIGH-v1.safetensors",
+    "WAN-2.2-I2V-Double-Blowjob-LOW-v1.safetensors",
+    "WAN-2.2-I2V-HandjobBlowjobCombo-HIGH-v1.safetensors",
+    "WAN-2.2-I2V-HandjobBlowjobCombo-LOW-v1.safetensors",
+    "WAN-2.2-I2V-SensualTeasingBlowjob-HIGH-v1.safetensors",
+    "WAN-2.2-I2V-SensualTeasingBlowjob-LOW-v1.safetensors",
+    "iGOON_Blink_Blowjob_I2V_HIGH.safetensors",
+    "iGOON_Blink_Blowjob_I2V_LOW.safetensors",
+    "iGoon - Blink_Front_Doggystyle_I2V_HIGH.safetensors",
+    "iGoon - Blink_Front_Doggystyle_I2V_LOW.safetensors",
+    "iGoon - Blink_Missionary_I2V_HIGH.safetensors",
+    "iGoon - Blink_Missionary_I2V_LOW v2.safetensors",
+    "iGoon - Blink_Missionary_I2V_LOW.safetensors",
+    "iGoon%20-%20Blink_Back_Doggystyle_HIGH.safetensors",
+    "iGoon%20-%20Blink_Back_Doggystyle_LOW.safetensors",
+    "iGoon%20-%20Blink_Facial_I2V_HIGH.safetensors",
+    "iGoon%20-%20Blink_Facial_I2V_LOW.safetensors",
+    "iGoon_Blink_Missionary_I2V_HIGH v2.safetensors",
+    "iGoon_Blink_Titjob_I2V_HIGH.safetensors",
+    "iGoon_Blink_Titjob_I2V_LOW.safetensors",
+    "lips-bj_high_noise.safetensors",
+    "lips-bj_low_noise.safetensors",
+    "mql_casting_sex_doggy_kneel_diagonally_behind_vagina_wan22_i2v_v1_high_noise.safetensors",
+    "mql_casting_sex_doggy_kneel_diagonally_behind_vagina_wan22_i2v_v1_low_noise.safetensors",
+    "mql_casting_sex_reverse_cowgirl_lie_front_vagina_wan22_i2v_v1_high_noise.safetensors",
+    "mql_casting_sex_reverse_cowgirl_lie_front_vagina_wan22_i2v_v1_low_noise.safetensors",
+    "mql_casting_sex_spoon_wan22_i2v_v1_high_noise.safetensors",
+    "mql_casting_sex_spoon_wan22_i2v_v1_low_noise.safetensors",
+    "mql_massage_tits_wan22_i2v_v1_high_noise.safetensors",
+    "mql_massage_tits_wan22_i2v_v1_low_noise.safetensors",
+    "mql_panties_aside_wan22_i2v_v1_high_noise.safetensors",
+    "mql_panties_aside_wan22_i2v_v1_low_noise.safetensors",
+    "sfbehind_v2.1_high_noise.safetensors",
+    "sfbehind_v2.1_low_noise.safetensors",
+    "sid3l3g_transition_v2.0_H.safetensors",
+    "sid3l3g_transition_v2.0_L.safetensors",
+    "wan2.2_i2v_high_ulitmate_pussy_asshole.safetensors",
+    "wan2.2_i2v_low_ulitmate_pussy_asshole.safetensors",
+    "wan22-mouthfull-140epoc-high-k3nk.safetensors",
+    "wan22-mouthfull-152epoc-low-k3nk.safetensors",
+]
+
+# group -> {"HIGH": (repo, file) | None, "LOW": (repo, file) | None}
+LORA_PAIRS = {}
+for f in LORA_FILES:
+    name = urllib.parse.unquote(f).replace(".safetensors", "")
+    is_high = bool(re.search(r'(high|HN|_H\b)', name, re.IGNORECASE))
+    is_low = bool(re.search(r'(low|LN|_L\b)', name, re.IGNORECASE))
+    group = re.sub(r'[\s_-]*(high|low|noise|HN|LN)([\s_-]*noise)?[\s_-]*(v?\d+(\.\d+)?)?\s*$', '', name, flags=re.IGNORECASE).strip()
+    group = re.sub(r'[\s_]+$', '', group)
+    LORA_PAIRS.setdefault(group, {"HIGH": None, "LOW": None})
+    if is_high:
+        LORA_PAIRS[group]["HIGH"] = (LORA_REPO, f)
+    elif is_low:
+        LORA_PAIRS[group]["LOW"] = (LORA_REPO, f)
+
+L = "Lythiga/WAN2.2_I2V_Lora"
+R = "rolechar/rc_wan2.2_i2v_loraset_nsfw_private"
+P = "lopi999/Wan2.2-I2V_General-NSFW-LoRA"
+
+# curated extra act/helper LoRAs (label -> (repo, high_file, low_file|None))
+EXTRA = {
+    # --- deepthroat ---
+    "Ultimate Deepthroat": (L, "wan22-ultimatedeepthroat-i2v-102epoc-high-k3nk.safetensors", "wan22-ultimatedeepthroat-I2V-101epoc-low-k3nk.safetensors"),
+    "BBC Deepthroat": (L, "wan22-bbcdeepthroat-115epoc-high-k3nk.safetensors", "wan22-bbcdeepthroat-155epoc-low-720-k3nk.safetensors"),
+    "Throat V2": (L, "Wan22_ThroatV2_High.safetensors", "Wan22_ThroatV2_Low.safetensors"),
+    "JFJ Deepthroat": (L, "jfj-deepthroat-W22-I2V-HN.safetensors", None),
+    # --- blowjob extra ---
+    "BBC Blowjob Extreme": (L, "BBC Blowjobs Extreme_high_noise.safetensors", "BBC Blowjobs Extreme_low_noise.safetensors"),
+    "POV Blowjob": (L, "Pov_Blowjob_Wan2.2_high_I2V_v1.0.safetensors", None),
+    # --- cunnilingus ---
+    "Cunnilingus": (L, "wan22-cunilingus-I2V-106epoc-high.safetensors", "wan22-cunilingus-I2V-72epoc-low.safetensors"),
+    # --- cowgirl ---
+    "Reverse Cowgirl v2": (L, "wan22.r3v3rs3_c0wg1rl-14b-High-i2v_e70.safetensors", "wan22.r3v3rs3_c0wg1rl-14b-Low-i2v_e70.safetensors"),
+    # --- doggy / from behind ---
+    "Doggy Front View v2": (L, "doggy_style_sex_front_view_Wan2.2_I2V_high_v1.0.safetensors", "doggy_style_sex_front_view_Wan2.2_I2V_Low_v1.0.safetensors"),
+    "Fucked From Behind": (L, "derpfuckedfrombehind_000002125_high_noise.safetensors", "derpfuckedfrombehind_000002125_low_noise.safetensors"),
+    "Prone Bone": (L, "Pronebone_high_noise.safetensors", "Pronebone_low_noise.safetensors"),
+    "Reverse Suspended Congress": (L, "reverse_suspended_congress_I2V_high.safetensors", "reverse_suspended_congress_I2V_low.safetensors"),
+    # --- missionary ---
+    "POV Missionary": (L, "wan2.2_i2v_highnoise_pov_missionary_v1.0.safetensors", "wan2.2_i2v_lownoise_pov_missionary_v1.0.safetensors"),
+    "Mating Press": (L, "mating_press_high.safetensors", "mating_press_low.safetensors"),
+    # --- standing ---
+    "Standing Upright Sex": (L, "Upright_sex_Wan2.2_14B_TI2V_HIGH_v1.0.safetensors", "Upright_sex_Wan2.2_14B_TI2V_LOW_v1.0.safetensors"),
+    # --- anal / ass ---
+    "Oral Insertion": (L, "wan2.2-i2v-high-oral-insertion-v1.0.safetensors", "wan2.2-i2v-low-oral-insertion-v1.0.safetensors"),
+    # --- handjob / balls ---
+    "Handjob": (L, "WAN-2.2-I2V-Handjob-HIGH-v1.safetensors", "WAN-2.2-I2V-Handjob-LOW-v1.safetensors"),
+    "Balls Sucking": (L, "WAN-2.2-I2V_Balls_sucking_High_noise.safetensors", "WAN-2.2-I2V_Balls_sucking_Low_noise.safetensors"),
+    # --- fingering ---
+    "Fingering": (R, "Sensual_fingering_v1_high_noise.safetensors", "Sensual_fingering_v1_low_noise.safetensors"),
+    # --- cumshot / creampie ---
+    "Facesplash Cumshot": (L, "wan22-f4c3spl4sh-100epoc-high-k3nk.safetensors", "wan22-f4c3spl4sh-154epoc-low-k3nk.safetensors"),
+    "Creampie CRM": (L, "Creampie CRM-FULL-EPOCH-80-HIGH.safetensors", "Creampie CRM-FULL-EPOCH-80-LOW.safetensors"),
+    "Pornmaster Creampie": (L, "Pornmaster_wan 2.2_14b_I2V_Creampie_v1_high_noise.safetensors", "Pornmaster_wan 2.2_14b_I2V_Creampie_v1_low_noise.safetensors"),
+    "Cum (generic)": (L, "Wan22_Cum_high_noise_1.V1.safetensors", "Wan22_Cum_low_noise_1.V1.safetensors"),
+    # --- fetish / expression ---
+    "French Kiss": (L, "WAN2.2-FrenchKiss_HighNoise.safetensors", "WAN2.2-FrenchKiss_LowNoise.safetensors"),
+    "Twerking": (L, "Twerking_I2V_high_noise.safetensors", "Twerking_I2V_low_noise.safetensors"),
+    "Ahegao Drooling": (R, "middlefinger_drooling_ahegao_high.safetensors", "middlefinger_drooling_ahegao_low.safetensors"),
+    # --- helpers: size / penetration / anatomy / camera / booster ---
+    "General NSFW Booster": (P, "NSFW-22-H-e8.safetensors", "NSFW-22-L-e8.safetensors"),
+    "Penetration Insert": (L, "PenInsert_high_noise.safetensors", "PenInsert_low_noise.safetensors"),
+    "Biggest Cock (size)": (L, "WAN-2.2-I2V_BiggestCock_high_noise_V1.safetensors", "WAN-2.2-I2V_BiggestCock_low_noise_V1.safetensors"),
+    "Doggy Slider": (L, "I2V_doggyslider_high.safetensors", "I2V_doggyslider_low.safetensors"),
+    "Pussy Helper": (R, "wan2.2_i2v_ulitmate_pussy_helper_high.safetensors", "wan2.2_i2v_ulitmate_pussy_helper_low.safetensors"),
+    "FOV Slider (camera)": (L, "wan2.2-i2v-high-sex-fov-slider-v1.0.safetensors", "wan2.2-i2v-low-sex-fov-slider-v1.0.safetensors"),
+    "Smashcut (camera)": (R, "wan2.2-i2v-sex-smashcut-v1.0-high.safetensors", "wan2.2-i2v-sex-smashcut-v1.0-low.safetensors"),
+}
+for label, (repo, hi, lo) in EXTRA.items():
+    LORA_PAIRS.setdefault(label, {"HIGH": None, "LOW": None})
+    if hi:
+        LORA_PAIRS[label]["HIGH"] = (repo, hi)
+    if lo:
+        LORA_PAIRS[label]["LOW"] = (repo, lo)
+
+
+def get_lora_choices():
+    choices = []
+    for group in sorted(LORA_PAIRS.keys()):
+        p = LORA_PAIRS[group]
+        if p["HIGH"] and p["LOW"]:
+            choices.append(group)
+        elif p["HIGH"]:
+            choices.append(f"{group} (HIGH only)")
+        elif p["LOW"]:
+            choices.append(f"{group} (LOW only)")
+    return choices
+
+
+def download_lora(group_name):
+    if not group_name:
+        return None, None
+    clean_name = re.sub(r'\s*\(HIGH only\)|\s*\(LOW only\)', '', group_name)
+    if clean_name not in LORA_PAIRS:
+        return None, None
+    pair = LORA_PAIRS[clean_name]
+    high_path, low_path = None, None
+    if pair["HIGH"]:
+        repo, fn = pair["HIGH"]
+        high_path = hf_hub_download(repo, fn, token=HF_TOKEN, revision=PINNED_REVISIONS.get(repo))
+    if pair["LOW"]:
+        repo, fn = pair["LOW"]
+        low_path = hf_hub_download(repo, fn, token=HF_TOKEN, revision=PINNED_REVISIONS.get(repo))
+    return high_path, low_path
+
+
+def load_lora_to_pipe(pipe, group_name, adapter_name="lora"):
+    high_path, low_path = download_lora(group_name)
+    if high_path and low_path:
+        pipe.load_lora_weights(high_path, adapter_name=f"{adapter_name}_high")
+        pipe.load_lora_weights(low_path, adapter_name=f"{adapter_name}_low")
+        print(f"Loaded LoRA pair: {group_name}")
+        return True
+    elif high_path:
+        pipe.load_lora_weights(high_path, adapter_name=adapter_name)
+        print(f"Loaded LoRA: {group_name}")
+        return True
+    elif low_path:
+        pipe.load_lora_weights(low_path, adapter_name=adapter_name)
+        print(f"Loaded LoRA (low): {group_name}")
+        return True
+    return False
+
+
+def unload_lora(pipe):
+    try:
+        pipe.unload_lora_weights()
+    except:
+        pass

model/loss.py

@@ -0,0 +1,128 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class EPE(nn.Module):
+    def __init__(self):
+        super(EPE, self).__init__()
+
+    def forward(self, flow, gt, loss_mask):
+        loss_map = (flow - gt.detach()) ** 2
+        loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5
+        return (loss_map * loss_mask)
+
+
+class Ternary(nn.Module):
+    def __init__(self):
+        super(Ternary, self).__init__()
+        patch_size = 7
+        out_channels = patch_size * patch_size
+        self.w = np.eye(out_channels).reshape(
+            (patch_size, patch_size, 1, out_channels))
+        self.w = np.transpose(self.w, (3, 2, 0, 1))
+        self.w = torch.tensor(self.w).float().to(device)
+
+    def transform(self, img):
+        patches = F.conv2d(img, self.w, padding=3, bias=None)
+        transf = patches - img
+        transf_norm = transf / torch.sqrt(0.81 + transf**2)
+        return transf_norm
+
+    def rgb2gray(self, rgb):
+        r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
+        gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
+        return gray
+
+    def hamming(self, t1, t2):
+        dist = (t1 - t2) ** 2
+        dist_norm = torch.mean(dist / (0.1 + dist), 1, True)
+        return dist_norm
+
+    def valid_mask(self, t, padding):
+        n, _, h, w = t.size()
+        inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t)
+        mask = F.pad(inner, [padding] * 4)
+        return mask
+
+    def forward(self, img0, img1):
+        img0 = self.transform(self.rgb2gray(img0))
+        img1 = self.transform(self.rgb2gray(img1))
+        return self.hamming(img0, img1) * self.valid_mask(img0, 1)
+
+
+class SOBEL(nn.Module):
+    def __init__(self):
+        super(SOBEL, self).__init__()
+        self.kernelX = torch.tensor([
+            [1, 0, -1],
+            [2, 0, -2],
+            [1, 0, -1],
+        ]).float()
+        self.kernelY = self.kernelX.clone().T
+        self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device)
+        self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device)
+
+    def forward(self, pred, gt):
+        N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3]
+        img_stack = torch.cat(
+            [pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0)
+        sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1)
+        sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1)
+        pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:]
+        pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:]
+
+        L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y)
+        loss = (L1X+L1Y)
+        return loss
+
+class MeanShift(nn.Conv2d):
+    def __init__(self, data_mean, data_std, data_range=1, norm=True):
+        c = len(data_mean)
+        super(MeanShift, self).__init__(c, c, kernel_size=1)
+        std = torch.Tensor(data_std)
+        self.weight.data = torch.eye(c).view(c, c, 1, 1)
+        if norm:
+            self.weight.data.div_(std.view(c, 1, 1, 1))
+            self.bias.data = -1 * data_range * torch.Tensor(data_mean)
+            self.bias.data.div_(std)
+        else:
+            self.weight.data.mul_(std.view(c, 1, 1, 1))
+            self.bias.data = data_range * torch.Tensor(data_mean)
+        self.requires_grad = False
+            
+class VGGPerceptualLoss(torch.nn.Module):
+    def __init__(self, rank=0):
+        super(VGGPerceptualLoss, self).__init__()
+        blocks = []
+        pretrained = True
+        self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features
+        self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, X, Y, indices=None):
+        X = self.normalize(X)
+        Y = self.normalize(Y)
+        indices = [2, 7, 12, 21, 30]
+        weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10/1.5]
+        k = 0
+        loss = 0
+        for i in range(indices[-1]):
+            X = self.vgg_pretrained_features[i](X)
+            Y = self.vgg_pretrained_features[i](Y)
+            if (i+1) in indices:
+                loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1
+                k += 1
+        return loss
+
+if __name__ == '__main__':
+    img0 = torch.zeros(3, 3, 256, 256).float().to(device)
+    img1 = torch.tensor(np.random.normal(
+        0, 1, (3, 3, 256, 256))).float().to(device)
+    ternary_loss = Ternary()
+    print(ternary_loss(img0, img1).shape)

model/pytorch_msssim/__init__.py

@@ -0,0 +1,198 @@
+import torch
+import torch.nn.functional as F
+from math import exp
+import numpy as np
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+    return gauss/gauss.sum()
+
+
+def create_window(window_size, channel=1):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device)
+    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+    return window
+
+def create_window_3d(window_size, channel=1):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t())
+    _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t())
+    window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device)
+    return window
+
+
+def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
+    # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
+    if val_range is None:
+        if torch.max(img1) > 128:
+            max_val = 255
+        else:
+            max_val = 1
+
+        if torch.min(img1) < -0.5:
+            min_val = -1
+        else:
+            min_val = 0
+        L = max_val - min_val
+    else:
+        L = val_range
+
+    padd = 0
+    (_, channel, height, width) = img1.size()
+    if window is None:
+        real_size = min(window_size, height, width)
+        window = create_window(real_size, channel=channel).to(img1.device).type_as(img1)
+    
+    mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel)
+    mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_mu2
+
+    C1 = (0.01 * L) ** 2
+    C2 = (0.03 * L) ** 2
+
+    v1 = 2.0 * sigma12 + C2
+    v2 = sigma1_sq + sigma2_sq + C2
+    cs = torch.mean(v1 / v2)  # contrast sensitivity
+
+    ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
+
+    if size_average:
+        ret = ssim_map.mean()
+    else:
+        ret = ssim_map.mean(1).mean(1).mean(1)
+
+    if full:
+        return ret, cs
+    return ret
+
+
+def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
+    # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
+    if val_range is None:
+        if torch.max(img1) > 128:
+            max_val = 255
+        else:
+            max_val = 1
+
+        if torch.min(img1) < -0.5:
+            min_val = -1
+        else:
+            min_val = 0
+        L = max_val - min_val
+    else:
+        L = val_range
+
+    padd = 0
+    (_, _, height, width) = img1.size()
+    if window is None:
+        real_size = min(window_size, height, width)
+        window = create_window_3d(real_size, channel=1).to(img1.device).type_as(img1)
+        # Channel is set to 1 since we consider color images as volumetric images
+
+    img1 = img1.unsqueeze(1)
+    img2 = img2.unsqueeze(1)
+
+    mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1)
+    mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_sq
+    sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu2_sq
+    sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_mu2
+
+    C1 = (0.01 * L) ** 2
+    C2 = (0.03 * L) ** 2
+
+    v1 = 2.0 * sigma12 + C2
+    v2 = sigma1_sq + sigma2_sq + C2
+    cs = torch.mean(v1 / v2)  # contrast sensitivity
+
+    ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
+
+    if size_average:
+        ret = ssim_map.mean()
+    else:
+        ret = ssim_map.mean(1).mean(1).mean(1)
+
+    if full:
+        return ret, cs
+    return ret
+
+
+def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False):
+    device = img1.device
+    weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device).type_as(img1)
+    levels = weights.size()[0]
+    mssim = []
+    mcs = []
+    for _ in range(levels):
+        sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range)
+        mssim.append(sim)
+        mcs.append(cs)
+
+        img1 = F.avg_pool2d(img1, (2, 2))
+        img2 = F.avg_pool2d(img2, (2, 2))
+
+    mssim = torch.stack(mssim)
+    mcs = torch.stack(mcs)
+
+    # Normalize (to avoid NaNs during training unstable models, not compliant with original definition)
+    if normalize:
+        mssim = (mssim + 1) / 2
+        mcs = (mcs + 1) / 2
+
+    pow1 = mcs ** weights
+    pow2 = mssim ** weights
+    # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/
+    output = torch.prod(pow1[:-1] * pow2[-1])
+    return output
+
+
+# Classes to re-use window
+class SSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True, val_range=None):
+        super(SSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.val_range = val_range
+
+        # Assume 3 channel for SSIM
+        self.channel = 3
+        self.window = create_window(window_size, channel=self.channel)
+
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.dtype == img1.dtype:
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype)
+            self.window = window
+            self.channel = channel
+
+        _ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average)
+        dssim = (1 - _ssim) / 2
+        return dssim
+
+class MSSSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True, channel=3):
+        super(MSSSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = channel
+
+    def forward(self, img1, img2):
+        return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average)

model/warplayer.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+backwarp_tenGrid = {}
+
+
+def warp(tenInput, tenFlow):
+    k = (str(tenFlow.device), str(tenFlow.size()))
+    if k not in backwarp_tenGrid:
+        tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=tenFlow.device).view(
+            1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
+        tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=tenFlow.device).view(
+            1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
+        backwarp_tenGrid[k] = torch.cat(
+            [tenHorizontal, tenVertical], 1).to(tenFlow.device)
+
+    tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
+                         tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1)
+
+    grid = backwarp_tenGrid[k].type_as(tenFlow)
+    
+    g = (grid + tenFlow).permute(0, 2, 3, 1)
+    return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True)

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,19 @@
+diffusers==0.38.0
+transformers==4.57.6
+accelerate===1.13.0
+safetensors
+sentencepiece
+peft==0.19.1
+ftfy
+imageio
+imageio-ffmpeg
+opencv-python
+torchao==0.17.0
+
+numpy>=1.16, <=1.23.5
+# tqdm>=4.35.0
+# sk-video>=1.1.10
+# opencv-python>=4.1.2
+# moviepy>=1.0.3
+torch==2.11.0
+torchvision==0.26.0