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_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,
    guidance_scale,
    guidance_scale_2,
    steps,
    seed,
    randomize_seed,
    progress,
):
    return steps * 15

@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,
    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)
            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
    ]
    generate_button.click(fn=generate_image, inputs=ui_inputs, outputs=[img_output, seed_input])

    

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