Spaces:
Running
Running
File size: 2,874 Bytes
812e69e 25d702c 812e69e ab613df 69a9a62 ab613df 2479e61 64d780d 812e69e 64d780d d60c82f 812e69e ec1fc83 64d780d ec1fc83 d60c82f 25d702c ec1fc83 812e69e da1c584 25761d6 28ae721 ec1fc83 78004f2 812e69e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import torch
from diffusers import UniPCMultistepScheduler
from diffusers import WanPipeline, AutoencoderKLWan # Use Wan-specific VAE
# from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
from diffusers.models import UNetSpatioTemporalConditionModel
from transformers import T5EncoderModel, T5Tokenizer
from PIL import Image
import numpy as np
import gradio as gr
import spaces
model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
# vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
# pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
flow_shift = 5.0 # 5.0 for 720P, 3.0 for 480P
# pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
@spaces.GPU()
def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, lora_id=None, progress=gr.Progress(track_tqdm=True)):
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
if lora_id and lora_id.strip() != "":
pipe.unload_lora_weights()
pipe.load_lora_weights(lora_id.strip())
pipe.to("cuda")
# apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
try:
output = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=1,
num_inference_steps=num_inference_steps,
guidance_scale=5.0,
)
image = output.frames[0][0]
image = (image * 255).astype(np.uint8)
return Image.fromarray(image)
finally:
if lora_id and lora_id.strip() != "":
pipe.unload_lora_weights()
iface = gr.Interface(
fn=generate,
inputs=[
gr.Textbox(label="Input prompt"),
],
additional_inputs = [
gr.Textbox(label="Negative prompt", value = "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"),
gr.Slider(label="Width", minimum=480, maximum=1280, step=16, value=1024),
gr.Slider(label="Height", minimum=480, maximum=1280, step=16, value=1024),
gr.Slider(minimum=1, maximum=80, step=1, label="Inference Steps", value=30),
gr.Textbox(label="LoRA ID"),
],
outputs=gr.Image(label="output"),
)
iface.launch() |