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

# Actual demo code
import spaces
import torch
from diffusers import WanPipeline, AutoencoderKLWan, UniPCMultistepScheduler
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_
import math


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

LANDSCAPE_WIDTH = 1024
LANDSCAPE_HEIGHT = 1024
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16

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('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')
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0)


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=81,
)



default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards, watermark, text, signature"


def get_duration(
    prompt,
    negative_prompt,
    guidance_scale,
    guidance_scale_2,
    steps,
    seed,
    randomize_seed,
    width,
    height,
    progress,
):
    megapixel = math.ceil((width * height) / 1_000_000)
    return steps * megapixel

@spaces.GPU(duration = get_duration)
def generate_image(
    prompt,
    negative_prompt=default_negative_prompt,
    guidance_scale = 1,
    guidance_scale_2 = 3,
    steps = 12,
    seed = 42,
    randomize_seed = False,
    width = 1024, 
    height = 1024,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate an image from a text prompt using the Wan 2.2 14B T2V model.
    
    This function takes an input prompt and generates an image based on the provided
    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Text-to-Video model.
    
    Args:
        prompt (str): Text prompt describing the desired image.
        negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
            Defaults to default_negative_prompt (contains unwanted visual artifacts).
        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:
            - image_path (str): Path to the generated image 
            - 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 uses GPU acceleration via the @spaces.GPU decorator
        - Generation time varies based on steps
    """
    
   
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    out_img = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_frames=1,
        guidance_scale=float(guidance_scale),
        guidance_scale_2=float(guidance_scale_2),
        num_inference_steps=int(steps),
        output_type="pil",
        generator=torch.Generator(device="cuda").manual_seed(current_seed),
    ).frames[0][0]

    return out_img, current_seed
css="""
#col-container {
    margin: 0 auto;
    max-width: 620px;
}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
       gr.Markdown("Fast Wan 2.2 T2I 14B")
       gr.Markdown("Generate high-quality images with Wan 2.2 14B and the Fast LoRa.")
       with gr.Row():
          prompt_input = gr.Textbox(show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,)
          generate_button = gr.Button("Run", variant="primary", scale=0)
       img_output = gr.Image(label="Generated Image", interactive=False)
       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=40, step=1, value=12, 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")
         
    ui_inputs = [ 
        prompt_input,
        negative_prompt_input,
        guidance_scale_input, guidance_scale_2_input, steps_slider, seed_input, randomize_seed_checkbox
    ]
    generate_button.click(fn=generate_image, inputs=ui_inputs, outputs=[img_output, seed_input])

    

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