Spaces:
Running on Zero
Running on Zero
File size: 3,183 Bytes
812e69e fed4c02 85f6fcb fed4c02 1b24a66 fed4c02 812e69e fed4c02 ab613df 69a9a62 ab613df fed4c02 5d27053 fed4c02 5d27053 adefe82 bf7c515 fed4c02 bf7c515 ac6ae94 fed4c02 ac6ae94 fed4c02 bf7c515 fed4c02 4166d00 09fcd4a c53b69e fed4c02 ac6ae94 fed4c02 d55351a fed4c02 ac6ae94 fed4c02 812e69e fed4c02 55713b6 fed4c02 812e69e fed4c02 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import torch
from diffusers import UniPCMultistepScheduler, FlowMatchEulerDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler
from diffusers import WanPipeline, AutoencoderKLWan # Use Wan-specific VAE
# from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
from diffusers.models import UNetSpatioTemporalConditionModel
from transformers import T5EncoderModel, T5Tokenizer
from huggingface_hub import hf_hub_download
from PIL import Image
import numpy as np
import gradio as gr
import spaces
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "Wan-AI/Wan2.1-T2V-14B-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 = 1.0 #5.0 1.0 for image, 5.0 for 720P, 3.0 for 480P
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
pipe.load_lora_weights(
"joerose/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32",
weight_name="Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors"
)
pipe.fuse_lora()
pipe.to(device)
# print("Initialization complete. Gradio is starting...")
@spaces.GPU(size='xlarge')
def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, lora_id=None, progress=gr.Progress(track_tqdm=True)):
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))
apply_cache_on_pipe(
pipe,
residual_diff_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=1.0, #5.0
)
image = output.frames[0][0]
image = (image * 255).astype(np.uint8)
return Image.fromarray(image)
finally:
pipe.unload_lora_weights()
# pass
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=2048, step=16, value=1024),
gr.Slider(label="Height", minimum=480, maximum=2048, step=16, value=1024),
gr.Slider(minimum=1, maximum=80, step=1, label="Inference Steps", value=8),
gr.Textbox(label="LoRA ID"),
],
outputs=gr.Image(label="output"),
)
iface.launch()
|