app1.py

import os
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')

# Actual demo code
import spaces
import torch
from diffusers import WanPipeline, AutoencoderKLWan
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from huggingface_hub import hf_hub_download, list_repo_files
from diffusers.utils.export_utils import export_to_video
import gradio as gr
import tempfile
import numpy as np
from PIL import Image
import random
import gc
from optimization import optimize_pipeline_

MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"

LANDSCAPE_WIDTH = 832
LANDSCAPE_HEIGHT = 480
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81

MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)

# Build the pipeline (bf16 on GPU)
vae = AutoencoderKLWan.from_pretrained(
    MODEL_ID, subfolder="vae", torch_dtype=torch.float32
)
pipe = WanPipeline.from_pretrained(
    MODEL_ID,
    transformer=WanTransformer3DModel.from_pretrained(
        "linoyts/Wan2.2-T2V-A14B-Diffusers-BF16",
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    ),
    transformer_2=WanTransformer3DModel.from_pretrained(
        "linoyts/Wan2.2-T2V-A14B-Diffusers-BF16",
        subfolder="transformer_2",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    ),
    vae=vae,
    torch_dtype=torch.bfloat16,
).to("cuda")

# ---- NEW: Load Orbit-Shot LoRA on the LOW-NOISE stage only (transformer_2) ----
# Repo: ostris/wan22_i2v_14b_orbit_shot_lora
# File: wan22_14b_i2v_orbit_low_noise.safetensors
#
# Notes:
# - We attach this LoRA only to transformer_2 (the low-noise stage for Wan2.2).
# - Adapter name `orbit_low` lets you adjust/enable later via set_adapters if needed.
# pipe.load_lora_weights(
#     "ostris/wan22_i2v_14b_orbit_shot_lora",
#     weight_name="wan22_14b_i2v_orbit_low_noise.safetensors",
#     adapter_name="orbit_low",
#     components=["transformer_2"],  # IMPORTANT: apply only to low-noise stage
# )
# # Activate the adapter on transformer_2 with weight 1.0 (changeable)
# pipe.set_adapters(
#     ["orbit_low"], adapter_weights=[1.0], components=["transformer_2"]
# )


LORA_REPO_ID = "deadman44/Wan2.2_T2i_T2v_LoRA"
LORA_FILENAME = "lora_wan2.2_myjd_Low_v01.safetensors"




    
# -------------------------------------------------------------------------------

# Usual CUDA cleanup
for _ in range(3):
    gc.collect()
    torch.cuda.synchronize()
    torch.cuda.empty_cache()

# Keep your current optimization hook
optimize_pipeline_(
    pipe,
    prompt="prompt",
    height=LANDSCAPE_HEIGHT,
    width=LANDSCAPE_WIDTH,
    num_frames=MAX_FRAMES_MODEL,
)

# ---- NEW: Load + Fuse Orbit-Shot LoRA on the LOW-NOISE stage only (transformer_2) ----
# Repo: ostris/wan22_i2v_14b_orbit_shot_lora
# File: wan22_14b_i2v_orbit_low_noise.safetensors
try:
    causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
    print("✅ LoRA downloaded to:", causvid_path)

    pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
    pipe.set_adapters(["causvid_lora"], adapter_weights=[0.75])
    pipe.fuse_lora()

except Exception as e:
    import traceback
    print("❌ Error during LoRA loading:")
    traceback.print_exc()
# -------------------------------------------------------------------------------


default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"

def get_duration(
    prompt,
    negative_prompt,
    duration_seconds,
    guidance_scale,
    guidance_scale_2,
    steps,
    seed,
    randomize_seed,
    progress,
):
    return steps * 15

@spaces.GPU(duration=get_duration)
def generate_video(
    prompt,
    negative_prompt=default_negative_prompt,
    duration_seconds=MAX_DURATION,
    guidance_scale=1,
    guidance_scale_2=3,
    steps=4,
    seed=42,
    randomize_seed=False,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate a video from a text prompt using the Wan 2.2 14B T2V model with a low-noise Orbit-Shot LoRA adapter.
    """
    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    output_frames_list = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=480,
        width=832,
        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),
    ).frames[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name

    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
    return video_path, current_seed

with gr.Blocks() as demo:
    gr.Markdown("# Fast 4 steps Wan 2.2 T2V (14B) + Orbit-Shot LoRA (low-noise)")
    gr.Markdown(
        "Runs Wan 2.2 in 4–8 steps. Low-noise **Orbit-Shot** LoRA is applied to `transformer_2` only."
    )
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v)
            duration_seconds_input = gr.Slider(
                minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION,
                label="Duration (seconds)",
                info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
            )
            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, 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")
                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=3, label="Guidance Scale 2 - low noise stage")
            generate_button = gr.Button("Generate Video", variant="primary")
        with gr.Column():
            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)

    ui_inputs = [
        prompt_input,
        negative_prompt_input,
        duration_seconds_input,
        guidance_scale_input,
        guidance_scale_2_input,
        steps_slider,
        seed_input,
        randomize_seed_checkbox,
    ]
    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])

    gr.Examples(
        examples=[
            [
                "POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind (clear water, green hills, blue sky with clouds). Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles and sunlight beams. Brief underwater view of cat’s face, then cat resurfaces, still filming selfie, playful summer vacation mood.",
            ],
            [
                "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
            ],
            [
                "A cinematic shot of a boat sailing on a calm sea at sunset.",
            ],
            [
                "Drone footage flying over a futuristic city with flying cars.",
            ],
        ],
        inputs=[prompt_input],
        outputs=[video_output, seed_input],
        fn=generate_video,
        cache_examples="lazy",
    )

if __name__ == "__main__":
    demo.queue().launch(mcp_server=True)