Hugging Face's logo

Spaces:
1mviews3
/
Wan2-1-Image-Generator
Running on Zero

App Files Community
Wan2-1-Image-Generator
File size: 4,397 Bytes
5024bca
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
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.to(device)



# Configure DDIMScheduler with a beta schedule
# pipe.scheduler = DDIMScheduler.from_config(
#     pipe.scheduler.config,
#     beta_start=0.00085,  # Starting beta value
#     beta_end=0.012,      # Ending beta value
#     beta_schedule="linear",  # Linear beta schedule (other options: "scaled_linear", "squaredcos_cap_v2")
#     num_train_timesteps=1000,  # Number of timesteps
#     flow_shift=flow_shift
# )






# Configure FlowMatchEulerDiscreteScheduler
# pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
#     pipe.scheduler.config,
#     flow_shift=flow_shift           # Retain flow_shift for WanPipeline compatibility
# )



# --- LoRA State Management ---
# Define unique names for our adapters
DEFAULT_LORA_NAME = "causvid_lora"
CUSTOM_LORA_NAME = "custom_lora"
# Track which custom LoRA is currently loaded to avoid reloading
CURRENTLY_LOADED_CUSTOM_LORA = None



# Load the default base LoRA ONCE at startup
print("Loading base LoRA...")
CAUSVID_LORA_REPO = "Kijai/WanVideo_comfy"
CAUSVID_LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32_v2.safetensors"
try:
   causvid_path = hf_hub_download(repo_id=CAUSVID_LORA_REPO, filename=CAUSVID_LORA_FILENAME)
   pipe.load_lora_weights(causvid_path, adapter_name=DEFAULT_LORA_NAME)
   print(f"✅ Default LoRA '{DEFAULT_LORA_NAME}' loaded successfully.")
except Exception as e:
   print(f"⚠️ Default LoRA could not be loaded: {e}")
   DEFAULT_LORA_NAME = None



# print("Initialization complete. Gradio is starting...")









@spaces.GPU()
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:
       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=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=10),
       gr.Textbox(label="LoRA ID"),
   ],
   outputs=gr.Image(label="output"),
)



iface.launch()