pipeline_manager.py

import spaces
import os
import sys
import copy
import time
import uuid
import tempfile
import torch
import torch._dynamo
import gradio as gr
from tqdm import tqdm
from huggingface_hub import HfApi
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 config
import aoti
import lora_loader
from image_utils import resize_image, resize_and_crop_to_match, get_num_frames
from rife_interp import rife_model, interpolate_bits, create_classic_boomerang_loop, create_ending_boomerang_loop, create_dynamic_boomerang_loop, create_adaptive_speed_ramping, call_sulphur_rife_api, clear_vram, is_cuda_usable
from face_swapper import swap_face_in_frames, swap_face_in_single_image
from prompt_relay import PromptRelayManager

pipe = WanImageToVideoPipeline.from_pretrained(
    config.MODEL_ID,
    torch_dtype=torch.bfloat16,
).to('cuda')
original_scheduler = copy.deepcopy(pipe.scheduler)

for i, lora in enumerate(config.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(config.LORA_MODELS)}") 
    except Exception as e:
        print("Error:", str(e))
        print("Failed LoRA:", name_high_tr)
        pipe.unload_lora_weights()

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')

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=False, lora_groups=None, 
    custom_lora_url="", custom_lora_scale=1.0, enable_prompt_relay=False,
    relay_prompt_schedule="", noise_temperature=1.0, extra_gpu_buffer=0, *args, **kwargs
):
    width, height = resized_image.size
    # Non-linear 3D attention memory & sequence scaling for Wan 2.2 frame count
    frame_ratio = (num_frames / 81.0) ** 1.38
    spatial_ratio = (width * height) / (832 * 624)
    factor = frame_ratio * spatial_ratio

    # Calibrated base step duration: 9.8s for <=4.0s (65 frames) to yield ~21s reservation (saving quota while covering ~18.5-19.2s GPU compute)
    BASE_STEP_DURATION = 9.8 if num_frames <= 65 else 11.0
    step_duration = BASE_STEP_DURATION * factor
    gen_time = int(steps) * step_duration

    # Automatically double reservation duration when Classifier-Free Guidance (GS > 1.0) is active
    if float(guidance_scale) > 1.0 or float(guidance_scale_2) > 1.0:
        gen_time = gen_time * 2.0

    overhead = 2.0 if num_frames <= 33 else (3.0 if num_frames <= 65 else 5.0)

    # Automatically add +3 seconds overhead if custom LoRA is requested
    cached_l = str(kwargs.get("cached_lora") or "").strip()
    if (custom_lora_url and str(custom_lora_url).strip()) or (cached_l and cached_l != "(None / Disable)"):
        overhead += 3.0

    total_time = overhead + gen_time + float(extra_gpu_buffer or 0)
    if safe_mode:
        total_time = total_time * 1.25

    return max(6, int(total_time) + 1)

@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, safe_mode=False, lora_groups=None, 
    custom_lora_url="", custom_lora_scale=1.0, enable_prompt_relay=False,
    relay_prompt_schedule="", noise_temperature=1.0, extra_gpu_buffer=0, progress=gr.Progress(track_tqdm=True)
):
    scheduler_class = config.SCHEDULER_MAP.get(scheduler_name)
    if scheduler_class.__name__ != pipe.scheduler.config._class_name or flow_shift != pipe.scheduler.config.get("flow_shift", "shift"):
        cfg = copy.deepcopy(original_scheduler.config)
        if scheduler_class.__name__ == "FlowMatchEulerDiscreteScheduler":
            cfg['shift'] = flow_shift
        else:
            cfg['flow_shift'] = flow_shift
        pipe.scheduler = scheduler_class.from_config(cfg)

    clear_vram()

    # Prompt Relay: Multi-Event Temporal Routing
    active_prompt = prompt
    if enable_prompt_relay and relay_prompt_schedule and str(relay_prompt_schedule).strip():
        events = PromptRelayManager.parse_schedule(relay_prompt_schedule, duration_seconds, num_frames)
        if events:
            print(f"🎬 Prompt Relay Active: {len(events)} temporal events routed across {duration_seconds}s")
            event_texts = [f"[{e['start_sec']}s-{e['end_sec']}s]: {e['prompt']}" for e in events]
            active_prompt = " ".join([e['prompt'] for e in events]) + " " + prompt
            print(f"   Combined Relay Prompt: {active_prompt[:100]}...")

    task_name = str(uuid.uuid4())[:8]
    print(f"Generating {num_frames} frames, task: {task_name}, {duration_seconds}, {resized_image.size}, lora={lora_groups}, custom_url={custom_lora_url}, temp={noise_temperature}")
    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}")

    if custom_lora_url and str(custom_lora_url).strip():
        try:
            loaded_custom = lora_loader.load_custom_url_lora(
                pipe, str(custom_lora_url).strip(), adapter_name="custom_civitai_lora", scale=float(custom_lora_scale)
            )
            if loaded_custom:
                lora_loaded = True
            else:
                gr.Warning("⚠️ Selected LoRA model is incompatible with Wan 2.2! File was automatically purged and video generation reverted to base model.")
        except Exception as e:
            print(f"Custom LoRA URL error: {e}")
            gr.Warning(f"⚠️ Custom LoRA Error: {e}. Video generation reverted to base Wan 2.2 model.")

    # Initial Noise Temperature scaling (0 Extra GPU Quota)
    latents = None
    if float(noise_temperature) != 1.0:
        try:
            latent_frames = (num_frames - 1) // 4 + 1
            latent_h = resized_image.height // 8
            latent_w = resized_image.width // 8
            gen = torch.Generator(device="cuda").manual_seed(current_seed)
            latents = torch.randn(
                (1, 16, latent_frames, latent_h, latent_w), 
                generator=gen, 
                device="cuda", 
                dtype=pipe.transformer.dtype
            ) * float(noise_temperature)
            print(f"🌡️ Noise Temperature applied: {noise_temperature} (latents scaled)")
        except Exception as e:
            print(f"Noise Temperature notice: {e}")
            latents = None

    pipe_kwargs = {
        "image": resized_image,
        "last_image": processed_last_image,
        "prompt": active_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 latents is not None:
        pipe_kwargs["latents"] = latents

    result = pipe(**pipe_kwargs)

    if lora_loaded:
        lora_loader.unload_lora(pipe)

    print("gen time passed:", time.time() - start)
    gpu_time = round(time.time() - start, 2)

    raw_frames_np = result.frames[0]
    pipe.scheduler = original_scheduler

    del result
    clear_vram()

    return raw_frames_np, task_name, gpu_time

def generate_video(
    input_image, last_image, prompt, steps=4, negative_prompt=config.default_negative_prompt, 
    duration_seconds=config.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, motion_extension_mode="⚡ Real-Time RIFE Interpolation (32/64 FPS Ultra-Smooth)",
    safe_mode=False, custom_lora_url="", custom_lora_scale=1.0,
    enable_prompt_relay=False, relay_prompt_schedule="",
    ref_face_image=None, target_gender="Any / All Faces",
    play_result_video=True, custom_filename="", noise_temperature=1.0,
    enable_vip_rife=False, vip_rife_multiplier="2x", vip_rife_mode="High-FPS Motion Smoothness (FPS Boost)",
    vip_rife_upscale=True,
    vip_rife_enhance_face=False,
    vip_password="",
    cached_lora="",
    extra_gpu_buffer=0,
    request: gr.Request = None,
    progress=gr.Progress(track_tqdm=True)
):
    if input_image is None:
        raise gr.Error("Please upload an input image.")

    hf_user = "Guest"
    user_ip = "Unknown"
    if request is not None:
        try:
            if hasattr(request, "headers") and request.headers:
                user_ip = (
                    request.headers.get("x-forwarded-for") or 
                    request.headers.get("x-real-ip") or 
                    request.headers.get("cf-connecting-ip") or 
                    getattr(getattr(request, "client", None), "host", "Unknown")
                )
                if "," in user_ip:
                    user_ip = user_ip.split(",")[0].strip()

                hf_name = (
                    request.headers.get("x-hf-user-name") or 
                    request.headers.get("x-hf-user") or 
                    request.headers.get("x-username") or 
                    getattr(request, "username", None)
                )
                if hf_name and str(hf_name).strip():
                    hf_user = str(hf_name).strip()
                elif user_ip and user_ip != "Unknown":
                    hf_user = f"Guest ({user_ip})"
                else:
                    hf_user = "Guest (Anonymous)"
            elif hasattr(request, "username") and request.username:
                hf_user = request.username
        except Exception as err:
            print(f"User identification notice: {err}")

    if hf_user == "Guest" and user_ip != "Unknown":
        hf_user = f"Guest ({user_ip})"

    active_custom_lora = str(custom_lora_url or "").strip()
    if not active_custom_lora and cached_lora and str(cached_lora).strip() != "(None / Disable)":
        active_custom_lora = str(cached_lora).strip()

    # CPU Pre-Download Custom LoRA (Before GPU inference starts to preserve ZeroGPU quota)
    if active_custom_lora:
        if active_custom_lora.startswith("http://") or active_custom_lora.startswith("https://"):
            start_dl = time.time()
            print(f"📥 Running CPU Pre-Download for Custom LoRA: {active_custom_lora}...")
            try:
                lora_path = lora_loader.download_file_from_url(active_custom_lora)
                print(f"✅ CPU Pre-Download complete in {time.time() - start_dl:.2f}s: {lora_path}")
            except Exception as e:
                print(f"❌ CPU Custom LoRA download failed: {e}")
                raise gr.Error(f"Gagal mengunduh LoRA dari URL: {e}")
        else:
            print(f"📦 Using Cached Custom LoRA: {active_custom_lora}")

    num_frames = get_num_frames(duration_seconds)
    current_seed = int(torch.randint(0, config.MAX_SEED, (1,)).item()) 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)

    reserved_time = get_inference_duration(
        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, None,
        active_custom_lora, custom_lora_scale, enable_prompt_relay,
        relay_prompt_schedule, noise_temperature, extra_gpu_buffer, progress, cached_lora=active_custom_lora
    )

    raw_frames_np, task_n, gpu_time = 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, None,
        active_custom_lora, custom_lora_scale, enable_prompt_relay,
        relay_prompt_schedule, noise_temperature, extra_gpu_buffer, progress
    )
    
    print(f"GPU complete: {task_n}. Release GPU lock and now processing post-processing on CPU...")

    # Motion Extension Technique & Playback FPS
    final_fps = config.FIXED_FPS
    mode_str = str(motion_extension_mode)

    if enable_vip_rife:
        print("💎 VIP RIFE Acceleration active: Bypassing local CPU post-processing...")
        final_frames = list(raw_frames_np)
        final_fps = config.FIXED_FPS
    elif "Ending" in mode_str or "Tail" in mode_str:
        print("🔂 Processing Ending-Only Boomerang Loop (Real-Speed Tail 1.5s Loop)...")
        final_frames = create_ending_boomerang_loop(raw_frames_np)
        final_fps = config.FIXED_FPS
    elif "Boomerang" in mode_str or "Loop" in mode_str or "Ping-Pong" in mode_str:
        print("🔂 Processing Classic Full Boomerang Loop (100% Real-Speed Forward+Reverse)...")
        final_frames = create_classic_boomerang_loop(raw_frames_np)
        final_fps = config.FIXED_FPS
    elif "Ramping" in mode_str or "Curve" in mode_str or "Ease" in mode_str or "Speed" in mode_str:
        print("🌊 Processing Adaptive Motion Speed Ramping (Ease-In/Out Curve, Real-Time Speed)...")
        final_frames = create_adaptive_speed_ramping(raw_frames_np, multiplier=2)
        final_fps = config.FIXED_FPS
    elif "Real-Time" in mode_str or "Ultra-Smooth" in mode_str:
        frame_factor = max(2, int(frame_multiplier // config.FIXED_FPS))
        calc_fps = int(frame_factor * config.FIXED_FPS)
        start = time.time()
        print(f"⚡ Processing Real-Time RIFE Interpolation ({calc_fps} FPS)...")
        use_cuda = is_cuda_usable()
        rife_device = torch.device("cuda" if use_cuda else "cpu")
        try:
            if use_cuda and hasattr(rife_model, "device"):
                rife_model.device()
                rife_model.flownet = rife_model.flownet.half()
            else:
                if hasattr(rife_model, "flownet") and rife_model.flownet is not None:
                    rife_model.flownet = rife_model.flownet.to(rife_device).float()
        except Exception as e:
            print(f"RIFE device setup notice: {e}")
        final_frames = interpolate_bits(raw_frames_np, multiplier=int(frame_factor))
        final_fps = calc_fps
        print("Interpolation time passed:", time.time() - start)
    else:
        # Classic Slow-Motion RIFE (16 FPS Time-Stretch)
        frame_factor = max(2, int(frame_multiplier // config.FIXED_FPS))
        start = time.time()
        print(f"🐢 Processing Slow-Motion RIFE Interpolation (16 FPS Time-Stretch, {frame_factor}x)...")
        use_cuda = is_cuda_usable()
        rife_device = torch.device("cuda" if use_cuda else "cpu")
        try:
            if use_cuda and hasattr(rife_model, "device"):
                rife_model.device()
                rife_model.flownet = rife_model.flownet.half()
            else:
                if hasattr(rife_model, "flownet") and rife_model.flownet is not None:
                    rife_model.flownet = rife_model.flownet.to(rife_device).float()
        except Exception as e:
            print(f"RIFE device setup notice: {e}")
        final_frames = interpolate_bits(raw_frames_np, multiplier=int(frame_factor))
        final_fps = config.FIXED_FPS
        print("Interpolation time passed:", time.time() - start)

    # Output Filename Logic
    if custom_filename and custom_filename.strip():
        filename = custom_filename.strip()
        if not filename.lower().endswith(".mp4"):
            filename += ".mp4"
        video_path = os.path.join(tempfile.gettempdir(), filename)
    else:
        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)

    # 💎 VIP Remote RIFE Acceleration (Mutually Exclusive: Bypasses CPU RIFE if Active & Authorized)
    if enable_vip_rife:
        vip_secret = (config.VIP_PASS or os.environ.get("VIP_PASSWORD", "")).strip()
        user_pass = (vip_password or "").strip()

        if vip_secret and user_pass != vip_secret:
            print("❌ Invalid VIP Password Access Key! Falling back to base output.")
            gr.Warning("❌ Invalid VIP Password Access Key! Remote VIP GPU RIFE acceleration was blocked.")
        else:
            try:
                print("💎 VIP Remote RIFE Acceleration authorized! Offloading to remote GPU engine...")
                mult_val = 2
                if "4x" in str(vip_rife_multiplier):
                    mult_val = 4
                elif "8x" in str(vip_rife_multiplier):
                    mult_val = 8

                is_slow_mo = ("Slow-Motion" in str(vip_rife_mode)) or ("Duration" in str(vip_rife_mode))
                
                vip_video_res = call_sulphur_rife_api(
                    video_path=video_path,
                    multiplier=mult_val,
                    slow_motion=is_slow_mo,
                    upscale=vip_rife_upscale,
                    enhance_face=vip_rife_enhance_face,
                    progress=progress
                )
                if vip_video_res and os.path.exists(vip_video_res):
                    video_path = vip_video_res
                    print(f"✅ VIP Remote RIFE Acceleration completed successfully: {video_path}")
                else:
                    print("⚠️ VIP RIFE Remote Acceleration failed or offline. Retaining base output.")
            except Exception as e:
                print(f"❌ VIP Remote RIFE error notice: {e}")

    # -------------------------------------------------------------
    # 💾 Instant Persistent Storage Auto-Save (/data/videos, /data/images, /data/prompts)
    # -------------------------------------------------------------
    storage_dir = "/data"
    if os.path.exists(storage_dir) and os.path.isdir(storage_dir):
        try:
            import shutil
            v_dir = os.path.join(storage_dir, "videos")
            i_dir = os.path.join(storage_dir, "images")
            p_dir = os.path.join(storage_dir, "prompts")

            os.makedirs(v_dir, exist_ok=True)
            os.makedirs(i_dir, exist_ok=True)
            os.makedirs(p_dir, exist_ok=True)

            v_filename = os.path.basename(video_path)
            v_basename = os.path.splitext(v_filename)[0]
            img_filename = f"input_{v_basename}.jpg"
            prompt_txt_filename = f"prompt_{v_basename}.txt"

            # 1. Save Video to /data/videos
            if os.path.exists(video_path):
                dest_v_path = os.path.join(v_dir, v_filename)
                shutil.copy(video_path, dest_v_path)
                print(f"💾 [Persistent Storage] Instantly saved video to: {dest_v_path}")

            # 2. Save Input Image to /data/images
            if input_image is not None:
                dest_i_path = os.path.join(i_dir, img_filename)
                input_image.convert("RGB").save(dest_i_path, format="JPEG", quality=95)
                print(f"💾 [Persistent Storage] Instantly saved input image to: {dest_i_path}")

            # 3. Save Prompt Text to /data/prompts
            prompt_body = f"user hf : {hf_user}\ngambar  : {img_filename}\nprompt  : {prompt}"
            if enable_prompt_relay and relay_prompt_schedule and str(relay_prompt_schedule).strip():
                prompt_body += f"\n\n[PROMPT RELAY SCHEDULE]\n{relay_prompt_schedule}"

            dest_p_path = os.path.join(p_dir, prompt_txt_filename)
            with open(dest_p_path, "w", encoding="utf-8") as f:
                f.write(prompt_body)
            print(f"💾 [Persistent Storage] Instantly saved prompt metadata to: {dest_p_path}")

        except Exception as e:
            print(f"⚠️ Persistent Storage Auto-Save notice : {e}")



    sec_per_step = round(gpu_time / max(1, int(steps)), 2)
    gpu_report_html = f"""
    <div style="background: rgba(99, 102, 241, 0.12); border: 1px solid rgba(99, 102, 241, 0.3); border-radius: 12px; padding: 12px 16px; margin-top: 12px; display: flex; align-items: center; justify-content: space-between; flex-wrap: wrap; gap: 8px;">
        <div style="display: flex; align-items: center; gap: 8px;">
            <span style="font-size: 1.1rem;">⚡</span>
            <span style="color: #a5b4fc; font-weight: 700; font-size: 0.92rem;">ZeroGPU Quota Consumed:</span>
            <span style="color: #38bdf8; font-weight: 800; font-size: 1.05rem;">{gpu_time:.2f} second</span>
        </div>
        <div style="display: flex; gap: 12px; font-size: 0.82rem; color: #94a3b8;">
            <span>Quota Reservation: <b>{reserved_time}s</b></span>
            <span>Speed: <b>{sec_per_step}s/step</b></span>
        </div>
    </div>
    """

    return (video_path if play_result_video else None), video_path, current_seed, gpu_report_html