app.py

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

# Actual demo code
import spaces
import torch
from diffusers import WanPipeline, AutoencoderKLWan
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
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)

vae = AutoencoderKLWan.from_pretrained("Wan-AI/Wan2.2-T2V-A14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)


# pipe = WanPipeline.from_pretrained(MODEL_ID,
#     transformer=WanTransformer3DModel.from_pretrained('rahul7star/wan2.2',
#         subfolder='Wan2.2-T2V-A14B-Diffusers-BF16/transformer',
#         torch_dtype=torch.bfloat16,
#         device_map='cuda',
#     ),
#     transformer_2=WanTransformer3DModel.from_pretrained('rahul7star/wan2.2',
#         subfolder='Wan2.2-T2V-A14B-Diffusers-BF16/transformer_2',
#         torch_dtype=torch.bfloat16,
#         device_map='cuda',
#     ),
#     vae=vae,
#     torch_dtype=torch.bfloat16,
# ).to('cuda')

HF_MODEL = os.environ.get("HF_UPLOAD_REPO", "rahul7star/WanText")
from huggingface_hub import HfApi, upload_file
import os
import uuid
import logging

import os
import uuid
import logging
from datetime import datetime
from huggingface_hub import HfApi, upload_file
import subprocess
import shutil
import tempfile
import logging
import subprocess
import tempfile
import logging
import shutil
import os
from huggingface_hub import HfApi, upload_file
from datetime import datetime
import uuid

import os
import tempfile
import logging
import subprocess
import uuid
import pandas as pd
from datetime import datetime
from huggingface_hub import upload_file, hf_hub_download

HF_MODEL = "rahul7star/WanText"         # replace with actual model repo
HF_DATASET_REPO = "rahul7star/Wan-video"


import os
import uuid
import tempfile
import logging
import subprocess
from datetime import datetime
import pandas as pd
from datasets import Dataset, Features, Value, Video
from huggingface_hub import upload_file, hf_hub_download


import os
import uuid
import tempfile
import logging
import subprocess
from datetime import datetime
import pandas as pd
from datasets import Dataset, Features, Value, Video
from huggingface_hub import upload_file, hf_hub_download



def upscale_and_upload_4k00(input_video_path: str, summary_text: str) -> str:
    """
    Upscale a video to 4K, upload to both model + dataset repos.
    Model repo: keeps videos + summary.txt in dated folders.
    Dataset repo: flat structure with growing metadata.csv (Video|Text|Date).
    Args:
        input_video_path (str): Path to original video.
        summary_text (str): Summary text.
    Returns:
        str: Hugging Face folder path inside model repo.
    """
    logging.info(f"Upscaling video to 4K for upload: {input_video_path}")

    # Temporary upscaled file
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_upscaled:
        upscaled_path = tmp_upscaled.name

    # FFmpeg upscale command
    cmd = [
        "ffmpeg",
        "-i", input_video_path,
        "-vf", "scale=3840:2160:flags=lanczos",
        "-c:v", "libx264",
        "-crf", "18",
        "-preset", "slow",
        "-y",
        upscaled_path,
    ]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        logging.info(f"✅ Upscaled video created at: {upscaled_path}")
    except subprocess.CalledProcessError as e:
        logging.error(f"FFmpeg failed:\n{e.stderr.decode()}")
        raise

    # Create date/unique folder for model repo
    today_str = datetime.now().strftime("%Y-%m-%d")
    unique_subfolder = f"Upload-4K-{uuid.uuid4().hex[:8]}"
    hf_folder = f"{today_str}/{unique_subfolder}"

    video_filename = os.path.basename(input_video_path)
    video_hf_path_model = f"{hf_folder}/{video_filename}"
    summary_hf_path_model = f"{hf_folder}/summary.txt"

    token = os.environ.get("HUGGINGFACE_HUB_TOKEN")

    def safe_upload(path, repo_id, repo_type, path_in_repo, label):
        try:
            upload_file(
                path_or_fileobj=path,
                path_in_repo=path_in_repo,
                repo_id=repo_id,
                repo_type=repo_type,
                token=token,
            )
            logging.info(f"✅ Uploaded {label} to {repo_type} repo {repo_id}: {path_in_repo}")
        except Exception as e:
            logging.error(f"❌ Failed to upload {label} to {repo_type} repo {repo_id}: {e}")

    # ----------------------
    # Upload to MODEL repo
    # ----------------------
    safe_upload(upscaled_path, HF_MODEL, "model", video_hf_path_model, "4K video")

    summary_file = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
    with open(summary_file, "w", encoding="utf-8") as f:
        f.write(summary_text)
    safe_upload(summary_file, HF_MODEL, "model", summary_hf_path_model, "summary")

    # ----------------------
    # Upload to DATASET repo (flat, growing metadata.csv)
    # ----------------------
    try:
        # Upload video to dataset repo root
        safe_upload(upscaled_path, HF_DATASET_REPO, "dataset", video_filename, "4K video")

        # Build absolute video URL for metadata
        video_url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/{video_filename}"

        # Load existing metadata.csv if exists
        rows = []
        try:
            existing_csv = hf_hub_download(
                repo_id=HF_DATASET_REPO,
                repo_type="dataset",
                filename="test.csv",
                token=token,
            )
            rows = pd.read_csv(existing_csv).to_dict("records")
        except Exception:
            logging.info("ℹ️ No existing metadata.csv found, creating new one.")

        # Append new row
        rows.append({"video": video_url, "text": summary_text, "date": today_str})
        df = pd.DataFrame(rows, columns=["video", "text", "date"])

        # Save and upload updated CSV
        csv_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
        df.to_csv(csv_path, index=False)
        safe_upload(csv_path, HF_DATASET_REPO, "dataset", "test.csv", "test.csv")

    except Exception as e:
        logging.error(f"❌ Dataset upload failed: {e}")

    # Cleanup temp files
    os.remove(upscaled_path)
    os.remove(summary_file)

    return hf_folder



def upscale_and_upload_4k(input_video_path: str, summary_text: str) -> str:
    """
    Upscale a video to 4K and upload it to Hugging Face Hub without replacing the original file.

    Args:
        input_video_path (str): Path to the original video.
        summary_text (str): Text summary to upload alongside the video.

    Returns:
        str: Hugging Face folder path where the video and summary were uploaded.
    """
    logging.info(f"Upscaling video to 4K for upload: {input_video_path}")

    # Create a temporary file for the upscaled video
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_upscaled:
        upscaled_path = tmp_upscaled.name

    # FFmpeg upscale command
    cmd = [
        "ffmpeg",
        "-i", input_video_path,
        "-vf", "scale=3840:2160:flags=lanczos",
        "-c:v", "libx264",
        "-crf", "18",
        "-preset", "slow",
        "-y",
        upscaled_path,
    ]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        logging.info(f"✅ Upscaled video created at: {upscaled_path}")
    except subprocess.CalledProcessError as e:
        logging.error(f"FFmpeg failed:\n{e.stderr.decode()}")
        raise

    # Create a date-based folder on HF
    today_str = datetime.now().strftime("%Y-%m-%d")
    unique_subfolder = f"Upload-4K-{uuid.uuid4().hex[:8]}"
    hf_folder = f"{today_str}/{unique_subfolder}"

    # Upload video
    video_filename = os.path.basename(input_video_path)
    video_hf_path = f"{hf_folder}/{video_filename}"
    upload_file(
        path_or_fileobj=upscaled_path,
        path_in_repo=video_hf_path,
        repo_id=HF_MODEL,
        repo_type="model",
        token=os.environ.get("HUGGINGFACE_HUB_TOKEN"),
    )
    logging.info(f"✅ Uploaded 4K video to HF: {video_hf_path}")

    # Upload summary.txt
    summary_file = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
    with open(summary_file, "w", encoding="utf-8") as f:
        f.write(summary_text)

    summary_hf_path = f"{hf_folder}/summary.txt"
    upload_file(
        path_or_fileobj=summary_file,
        path_in_repo=summary_hf_path,
        repo_id=HF_MODEL,
        repo_type="model",
        token=os.environ.get("HUGGINGFACE_HUB_TOKEN"),
    )
    logging.info(f"✅ Uploaded summary to HF: {summary_hf_path}")

    # Cleanup temporary files
    os.remove(upscaled_path)
    os.remove(summary_file)

    return hf_folder

def upload_to_hf(video_path, summary_text):
    api = HfApi()
    
    # Create a date-based folder (YYYY-MM-DD)
    today_str = datetime.now().strftime("%Y-%m-%d")
    date_folder = today_str
    
    # Generate a unique subfolder for this upload
    unique_subfolder = f"WanT2V-upload_{uuid.uuid4().hex[:8]}"
    hf_folder = f"{date_folder}/{unique_subfolder}"
    logging.info(f"Uploading files to HF folder: {hf_folder} in repo {HF_MODEL}")

    # Upload video
    video_filename = os.path.basename(video_path)
    video_hf_path = f"{hf_folder}/{video_filename}"
    upload_file(
        path_or_fileobj=video_path,
        path_in_repo=video_hf_path,
        repo_id=HF_MODEL,
        repo_type="model",
        token=os.environ.get("HUGGINGFACE_HUB_TOKEN"),
    )
    logging.info(f"✅ Uploaded video to HF: {video_hf_path}")

    # Upload summary.txt
    summary_file = "/tmp/summary.txt"
    with open(summary_file, "w", encoding="utf-8") as f:
        f.write(summary_text)

    summary_hf_path = f"{hf_folder}/summary.txt"
    upload_file(
        path_or_fileobj=summary_file,
        path_in_repo=summary_hf_path,
        repo_id=HF_MODEL,
        repo_type="model",
        token=os.environ.get("HUGGINGFACE_HUB_TOKEN"),
    )
    logging.info(f"✅ Uploaded summary to HF: {summary_hf_path}")

    return hf_folder



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


for i in range(3): 
    gc.collect()
    torch.cuda.synchronize() 
    torch.cuda.empty_cache()

optimize_pipeline_(pipe,
    prompt='prompt',
    height=LANDSCAPE_HEIGHT,
    width=LANDSCAPE_WIDTH,
    num_frames=MAX_FRAMES_MODEL,
)


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 Lightning LoRA.
    
    This function takes an input prompt and generates a video animation based on the provided
    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Text-to-Video model with Lightning LoRA
    for fast generation in 4-8 steps.
    
    Args:
        prompt (str): Text prompt describing the desired animation or motion.
        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.
        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
            Defaults to 4. Range: 1-30.
        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.
        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
    
    Returns:
        tuple: A tuple containing:
            - video_path (str): Path to the generated video file (.mp4)
            - current_seed (int): The seed used for generation (useful when randomize_seed=True)
    
    Raises:
        gr.Error: If input_image is None (no image uploaded).
    
    Note:
        - The function automatically resizes the input image to the target dimensions
        - 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)
    """
    
    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)
    upscale_and_upload_4k(video_path, prompt)
    return video_path, current_seed

with gr.Blocks() as demo:
    gr.Markdown("# Fast 4 steps Wan 2.2 T2V (14B) with Lightning LoRA")
    gr.Markdown("run Wan 2.2 in just 4-8 steps, with [Wan 2.2 Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU⚡️")
    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)