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_


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
MIN_FRAMES_MODEL = 1
MAX_FRAMES_MODEL = 81

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



default_negative_prompt = ""


def get_duration(
    prompt,
    negative_prompt,
    guidance_scale,
    guidance_scale_2,
    steps,
    seed,
    randomize_seed,
    # width,
    # height,
    progress,
):
    return steps * 2

@spaces.GPU(duration = get_duration)
def generate_image(
    prompt,
    negative_prompt=default_negative_prompt,
    guidance_scale = 3.5,
    guidance_scale_2 = 4,
    steps = 27,
    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=1024,
        width=1024,
        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("# Wan 2.2 (14B) Image")
       gr.Markdown("generate high quality images with Wan 2.2 14B")
       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)
            with gr.Row():
                width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1280, step=16)
                height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1280, step=16)
            steps_slider = gr.Slider(minimum=1, maximum=40, step=1, value=27, label="Inference Steps") 
            guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=3.5, label="Guidance Scale - high noise stage")
            guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=4, label="Guidance Scale 2 - low noise stage")

       gr.Examples(
        examples=[
            [
                "Extreme close-up portrait of a theatrical villain character with chalk-white face paint, blood-red lips stretched into an unsettling wide grin, wild green hair disheveled and unkempt, piercing eyes with dark smudged makeup, dramatic high contrast lighting casting deep shadows across angular facial features, menacing expression with visible teeth, gritty urban atmosphere, cinematic quality, sharp focus on facial details"
            ], 
            [
                "Whimsical fiber art sculpture of a wise green alien sage character crafted entirely from soft wool yarn, sitting contentedly in a vintage clawfoot bathtub filled with fluffy white soap bubbles, the yarn figure's earth-tone knitted robes darkened and soggy from bathwater, intricate crochet textures visible on the wet yarn surface, iridescent soap bubbles floating through the air catching rainbow reflections, small rubber duck and bath toys scattered around the tub, warm bathroom lighting creating cozy shadows, steam gently rising from the warm water, the character's embroidered facial features showing peaceful relaxation, macro photography highlighting the contrast between wet yarn fibers and delicate soap foam, quirky handcrafted art piece with playful domestic setting"
            ],
            [
                "Macro photograph of a honey bee delicately perched on vibrant purple lavender blossoms during golden hour, warm amber sunlight filtering through translucent wings, soft bokeh background with gradient sky transitioning from orange to deep blue, dewdrops glistening on flower petals, shallow depth of field emphasizing the bee's fuzzy texture and intricate wing details, peaceful twilight atmosphere, cinematic lighting, ultra-sharp focus on the bee"
            ],
        ],
        inputs=[prompt_input], outputs=[img_output, seed_input], fn=generate_image, cache_examples="lazy")            
    ui_inputs = [ 
        prompt_input,
        negative_prompt_input,
        guidance_scale_input, guidance_scale_2_input, steps_slider, seed_input, randomize_seed_checkbox 
        # width, height
    ]
    generate_button.click(fn=generate_image, inputs=ui_inputs, outputs=[img_output, seed_input])

    

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