feat: test 2nd pass
Browse files
app.py
CHANGED
|
@@ -5,6 +5,8 @@ from diffusers.utils import export_to_video
|
|
| 5 |
from diffusers import AutoencoderKLWan, WanPipeline
|
| 6 |
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
|
| 7 |
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
|
|
|
|
|
|
| 8 |
|
| 9 |
# Define model options
|
| 10 |
MODEL_OPTIONS = {
|
|
@@ -33,7 +35,14 @@ def generate_video(
|
|
| 33 |
num_frames,
|
| 34 |
guidance_scale,
|
| 35 |
num_inference_steps,
|
| 36 |
-
output_fps
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 37 |
):
|
| 38 |
# Get model ID from selection
|
| 39 |
model_id = MODEL_OPTIONS[model_choice]
|
|
@@ -76,35 +85,123 @@ def generate_video(
|
|
| 76 |
# Enable CPU offload for low VRAM
|
| 77 |
pipe.enable_model_cpu_offload()
|
| 78 |
|
| 79 |
-
#
|
| 80 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 81 |
prompt=prompt,
|
| 82 |
negative_prompt=negative_prompt,
|
| 83 |
height=height,
|
| 84 |
width=width,
|
| 85 |
num_frames=num_frames,
|
| 86 |
guidance_scale=guidance_scale,
|
| 87 |
-
num_inference_steps=num_inference_steps
|
| 88 |
-
|
|
|
|
|
|
|
| 89 |
|
| 90 |
-
#
|
| 91 |
-
|
| 92 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 93 |
|
| 94 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 95 |
|
| 96 |
# Create the Gradio interface
|
| 97 |
with gr.Blocks() as demo:
|
| 98 |
gr.HTML("""
|
| 99 |
<p align="center">
|
| 100 |
<svg version="1.1" viewBox="0 0 1200 295" xmlns="http://www.w3.org/2000/svg" xmlns:v="https://vecta.io/nano" width="400">
|
| 101 |
-
|
| 102 |
</svg>
|
| 103 |
<p align="center">
|
| 104 |
💻 <a href="https://www.markury.dev/"><b>Website</b></a>    |    🤗 <a href="https://huggingface.co/markury"><b>Hugging Face</b></a>    |    💿 <a href="https://thebulge.xyz"><b>Discord</b></a>
|
| 105 |
</p>
|
| 106 |
""")
|
| 107 |
-
gr.Markdown("# Wan 2.1 T2V 1.3B with LoRA")
|
| 108 |
|
| 109 |
with gr.Row():
|
| 110 |
with gr.Column(scale=1):
|
|
@@ -209,13 +306,92 @@ with gr.Blocks() as demo:
|
|
| 209 |
step=1
|
| 210 |
)
|
| 211 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 212 |
generate_btn = gr.Button("Generate Video")
|
| 213 |
|
| 214 |
with gr.Column(scale=1):
|
| 215 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 216 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 217 |
generate_btn.click(
|
| 218 |
-
fn=
|
| 219 |
inputs=[
|
| 220 |
model_choice,
|
| 221 |
prompt,
|
|
@@ -230,9 +406,20 @@ with gr.Blocks() as demo:
|
|
| 230 |
num_frames,
|
| 231 |
guidance_scale,
|
| 232 |
num_inference_steps,
|
| 233 |
-
output_fps
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 234 |
],
|
| 235 |
-
outputs=
|
|
|
|
|
|
|
|
|
|
|
|
|
| 236 |
)
|
| 237 |
|
| 238 |
gr.Markdown("""
|
|
@@ -242,6 +429,12 @@ with gr.Blocks() as demo:
|
|
| 242 |
- Number of frames should be of the form 4k+1 (e.g., 33, 81)
|
| 243 |
- Stick to lower frame counts. Even at 480p, an 81 frame sequence at 30 steps will nearly time out the request in this space.
|
| 244 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 245 |
## Using LoRAs with multiple safetensors files:
|
| 246 |
If you encounter an error stating "more than one weights file", you need to specify the exact weight file name in the "LoRA Weight Name" field.
|
| 247 |
You can find this by browsing the repository on Hugging Face and looking for the safetensors files (common names include: adapter_model.safetensors, pytorch_lora_weights.safetensors).
|
|
|
|
| 5 |
from diffusers import AutoencoderKLWan, WanPipeline
|
| 6 |
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
|
| 7 |
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
| 8 |
+
import os
|
| 9 |
+
import tempfile
|
| 10 |
|
| 11 |
# Define model options
|
| 12 |
MODEL_OPTIONS = {
|
|
|
|
| 35 |
num_frames,
|
| 36 |
guidance_scale,
|
| 37 |
num_inference_steps,
|
| 38 |
+
output_fps,
|
| 39 |
+
# Second pass parameters
|
| 40 |
+
enable_second_pass,
|
| 41 |
+
second_pass_scale,
|
| 42 |
+
second_pass_denoise,
|
| 43 |
+
second_pass_flow_shift,
|
| 44 |
+
second_pass_cfg,
|
| 45 |
+
show_both_outputs
|
| 46 |
):
|
| 47 |
# Get model ID from selection
|
| 48 |
model_id = MODEL_OPTIONS[model_choice]
|
|
|
|
| 85 |
# Enable CPU offload for low VRAM
|
| 86 |
pipe.enable_model_cpu_offload()
|
| 87 |
|
| 88 |
+
# Keep track of output files for return
|
| 89 |
+
output_files = []
|
| 90 |
+
|
| 91 |
+
# First pass - generate base video
|
| 92 |
+
print("Running first pass...")
|
| 93 |
+
first_pass = pipe(
|
| 94 |
prompt=prompt,
|
| 95 |
negative_prompt=negative_prompt,
|
| 96 |
height=height,
|
| 97 |
width=width,
|
| 98 |
num_frames=num_frames,
|
| 99 |
guidance_scale=guidance_scale,
|
| 100 |
+
num_inference_steps=num_inference_steps,
|
| 101 |
+
output_type="latent" if enable_second_pass else "pt", # Only return latents if doing second pass
|
| 102 |
+
return_dict=True
|
| 103 |
+
)
|
| 104 |
|
| 105 |
+
# Get the latents from the first pass output
|
| 106 |
+
latents = first_pass.frames[0]
|
| 107 |
+
|
| 108 |
+
# If we're not doing a second pass or need to display both outputs, decode the first pass
|
| 109 |
+
if not enable_second_pass or (enable_second_pass and show_both_outputs):
|
| 110 |
+
# Decode the latents to frames with the VAE (only needed if we requested latents)
|
| 111 |
+
if enable_second_pass:
|
| 112 |
+
print("Decoding first pass latents...")
|
| 113 |
+
with torch.no_grad():
|
| 114 |
+
first_pass_frames = pipe.vae.decode(latents / pipe.vae.config.scaling_factor).sample
|
| 115 |
+
else:
|
| 116 |
+
first_pass_frames = latents
|
| 117 |
+
|
| 118 |
+
# Export first pass to video
|
| 119 |
+
first_pass_file = "output_first_pass.mp4"
|
| 120 |
+
export_to_video(first_pass_frames, first_pass_file, fps=output_fps)
|
| 121 |
+
output_files.append(first_pass_file)
|
| 122 |
+
|
| 123 |
+
# Second pass - upscale and refine if enabled
|
| 124 |
+
second_pass_file = None
|
| 125 |
+
if enable_second_pass:
|
| 126 |
+
print("Running second pass with scale factor:", second_pass_scale)
|
| 127 |
+
|
| 128 |
+
# Resize latents for second pass (upscale)
|
| 129 |
+
new_height = int(height * second_pass_scale)
|
| 130 |
+
new_width = int(width * second_pass_scale)
|
| 131 |
+
|
| 132 |
+
# Ensure dimensions are multiples of 8
|
| 133 |
+
new_height = (new_height // 8) * 8
|
| 134 |
+
new_width = (new_width // 8) * 8
|
| 135 |
+
|
| 136 |
+
print(f"Upscaling latents from {height}x{width} to {new_height}x{new_width}")
|
| 137 |
+
|
| 138 |
+
# Upscale latents using interpolate
|
| 139 |
+
upscaled_latents = torch.nn.functional.interpolate(
|
| 140 |
+
latents,
|
| 141 |
+
size=(num_frames, new_height // 8, new_width // 8), # VAE downsamples by factor of 8
|
| 142 |
+
mode="trilinear",
|
| 143 |
+
align_corners=False
|
| 144 |
+
)
|
| 145 |
+
|
| 146 |
+
# Update scheduler for second pass if using different flow shift
|
| 147 |
+
if scheduler_type == "UniPCMultistepScheduler":
|
| 148 |
+
pipe.scheduler = UniPCMultistepScheduler.from_config(
|
| 149 |
+
pipe.scheduler.config,
|
| 150 |
+
flow_shift=second_pass_flow_shift
|
| 151 |
+
)
|
| 152 |
+
else:
|
| 153 |
+
pipe.scheduler = FlowMatchEulerDiscreteScheduler(shift=second_pass_flow_shift)
|
| 154 |
+
|
| 155 |
+
# Calculate noise level for partial denoising
|
| 156 |
+
# For noise scheduler, 0 means no noise (final step) and 1 means full noise (first step)
|
| 157 |
+
# So we convert our denoise strength to a timestep in the schedule
|
| 158 |
+
start_step = int(second_pass_denoise * num_inference_steps)
|
| 159 |
+
|
| 160 |
+
# Run second pass with the upscaled latents and partial denoising
|
| 161 |
+
print(f"Denoising from step {start_step} of {num_inference_steps} (denoise strength: {second_pass_denoise})")
|
| 162 |
+
|
| 163 |
+
# Use the second pass CFG value
|
| 164 |
+
second_pass_guidance = second_pass_cfg if second_pass_cfg > 0 else guidance_scale
|
| 165 |
+
|
| 166 |
+
second_pass = pipe(
|
| 167 |
+
prompt=prompt,
|
| 168 |
+
negative_prompt=negative_prompt,
|
| 169 |
+
height=new_height,
|
| 170 |
+
width=new_width,
|
| 171 |
+
num_frames=num_frames,
|
| 172 |
+
guidance_scale=second_pass_guidance,
|
| 173 |
+
num_inference_steps=num_inference_steps,
|
| 174 |
+
latents=upscaled_latents, # Use the upscaled latents
|
| 175 |
+
strength=second_pass_denoise, # Partial denoising
|
| 176 |
+
output_type="pt",
|
| 177 |
+
return_dict=True
|
| 178 |
+
)
|
| 179 |
+
|
| 180 |
+
# Export second pass to video
|
| 181 |
+
second_pass_file = "output_second_pass.mp4"
|
| 182 |
+
export_to_video(second_pass.frames[0], second_pass_file, fps=output_fps)
|
| 183 |
+
output_files.append(second_pass_file)
|
| 184 |
|
| 185 |
+
# Return the appropriate video output(s)
|
| 186 |
+
if enable_second_pass and not show_both_outputs:
|
| 187 |
+
return second_pass_file
|
| 188 |
+
elif enable_second_pass and show_both_outputs:
|
| 189 |
+
return [first_pass_file, second_pass_file]
|
| 190 |
+
else:
|
| 191 |
+
return first_pass_file
|
| 192 |
|
| 193 |
# Create the Gradio interface
|
| 194 |
with gr.Blocks() as demo:
|
| 195 |
gr.HTML("""
|
| 196 |
<p align="center">
|
| 197 |
<svg version="1.1" viewBox="0 0 1200 295" xmlns="http://www.w3.org/2000/svg" xmlns:v="https://vecta.io/nano" width="400">
|
| 198 |
+
...
|
| 199 |
</svg>
|
| 200 |
<p align="center">
|
| 201 |
💻 <a href="https://www.markury.dev/"><b>Website</b></a>    |    🤗 <a href="https://huggingface.co/markury"><b>Hugging Face</b></a>    |    💿 <a href="https://thebulge.xyz"><b>Discord</b></a>
|
| 202 |
</p>
|
| 203 |
""")
|
| 204 |
+
gr.Markdown("# Wan 2.1 T2V 1.3B with LoRA and Second Pass Refinement")
|
| 205 |
|
| 206 |
with gr.Row():
|
| 207 |
with gr.Column(scale=1):
|
|
|
|
| 306 |
step=1
|
| 307 |
)
|
| 308 |
|
| 309 |
+
# Add Second Pass options
|
| 310 |
+
with gr.Accordion("Second Pass Refinement (HiresFix)", open=False):
|
| 311 |
+
enable_second_pass = gr.Checkbox(
|
| 312 |
+
label="Enable Second Pass Refinement",
|
| 313 |
+
value=False,
|
| 314 |
+
info="Scale up and refine the video with a second denoising pass"
|
| 315 |
+
)
|
| 316 |
+
|
| 317 |
+
with gr.Row():
|
| 318 |
+
second_pass_scale = gr.Slider(
|
| 319 |
+
label="Scale Factor",
|
| 320 |
+
minimum=1.0,
|
| 321 |
+
maximum=2.0,
|
| 322 |
+
value=1.25,
|
| 323 |
+
step=0.05,
|
| 324 |
+
info="How much to upscale the video for refinement"
|
| 325 |
+
)
|
| 326 |
+
second_pass_denoise = gr.Slider(
|
| 327 |
+
label="Denoise Strength",
|
| 328 |
+
minimum=0.1,
|
| 329 |
+
maximum=1.0,
|
| 330 |
+
value=0.6,
|
| 331 |
+
step=0.05,
|
| 332 |
+
info="Lower values preserve more of the original details"
|
| 333 |
+
)
|
| 334 |
+
|
| 335 |
+
with gr.Row():
|
| 336 |
+
second_pass_flow_shift = gr.Slider(
|
| 337 |
+
label="Second Pass Flow Shift",
|
| 338 |
+
minimum=1.0,
|
| 339 |
+
maximum=12.0,
|
| 340 |
+
value=3.0,
|
| 341 |
+
step=0.5,
|
| 342 |
+
info="Flow shift value for the second pass (optional)"
|
| 343 |
+
)
|
| 344 |
+
second_pass_cfg = gr.Slider(
|
| 345 |
+
label="Second Pass CFG",
|
| 346 |
+
minimum=0.0,
|
| 347 |
+
maximum=15.0,
|
| 348 |
+
value=0.0,
|
| 349 |
+
step=0.5,
|
| 350 |
+
info="Set to 0 to use the same value as first pass"
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
show_both_outputs = gr.Checkbox(
|
| 354 |
+
label="Show Both Outputs",
|
| 355 |
+
value=False,
|
| 356 |
+
info="Display both original and refined videos"
|
| 357 |
+
)
|
| 358 |
+
|
| 359 |
generate_btn = gr.Button("Generate Video")
|
| 360 |
|
| 361 |
with gr.Column(scale=1):
|
| 362 |
+
# Updated output to handle multiple videos if both outputs are selected
|
| 363 |
+
with gr.Group():
|
| 364 |
+
output_video = gr.Video(label="Generated Video")
|
| 365 |
+
second_output_video = gr.Video(label="Second Pass Video", visible=False)
|
| 366 |
+
|
| 367 |
+
# Show/hide second video based on checkbox
|
| 368 |
+
def update_second_video_visibility(enable_pass, show_both):
|
| 369 |
+
return {"visible": enable_pass and show_both}
|
| 370 |
+
|
| 371 |
+
enable_second_pass.change(
|
| 372 |
+
fn=update_second_video_visibility,
|
| 373 |
+
inputs=[enable_second_pass, show_both_outputs],
|
| 374 |
+
outputs=[second_output_video]
|
| 375 |
+
)
|
| 376 |
+
|
| 377 |
+
show_both_outputs.change(
|
| 378 |
+
fn=update_second_video_visibility,
|
| 379 |
+
inputs=[enable_second_pass, show_both_outputs],
|
| 380 |
+
outputs=[second_output_video]
|
| 381 |
+
)
|
| 382 |
|
| 383 |
+
# Updated function to handle the second pass and multiple outputs
|
| 384 |
+
def process_generation(*args):
|
| 385 |
+
result = generate_video(*args)
|
| 386 |
+
if isinstance(result, list) and len(result) > 1:
|
| 387 |
+
return [result[0], result[1], {"visible": True}]
|
| 388 |
+
elif isinstance(result, list) and len(result) == 1:
|
| 389 |
+
return [result[0], None, {"visible": False}]
|
| 390 |
+
else:
|
| 391 |
+
return [result, None, {"visible": False}]
|
| 392 |
+
|
| 393 |
generate_btn.click(
|
| 394 |
+
fn=process_generation,
|
| 395 |
inputs=[
|
| 396 |
model_choice,
|
| 397 |
prompt,
|
|
|
|
| 406 |
num_frames,
|
| 407 |
guidance_scale,
|
| 408 |
num_inference_steps,
|
| 409 |
+
output_fps,
|
| 410 |
+
# Second pass parameters
|
| 411 |
+
enable_second_pass,
|
| 412 |
+
second_pass_scale,
|
| 413 |
+
second_pass_denoise,
|
| 414 |
+
second_pass_flow_shift,
|
| 415 |
+
second_pass_cfg,
|
| 416 |
+
show_both_outputs
|
| 417 |
],
|
| 418 |
+
outputs=[
|
| 419 |
+
output_video,
|
| 420 |
+
second_output_video,
|
| 421 |
+
second_output_video # Update visibility
|
| 422 |
+
]
|
| 423 |
)
|
| 424 |
|
| 425 |
gr.Markdown("""
|
|
|
|
| 429 |
- Number of frames should be of the form 4k+1 (e.g., 33, 81)
|
| 430 |
- Stick to lower frame counts. Even at 480p, an 81 frame sequence at 30 steps will nearly time out the request in this space.
|
| 431 |
|
| 432 |
+
## Second Pass Refinement Tips:
|
| 433 |
+
- The second pass (similar to HiresFix) can enhance details by upscaling and refining the video
|
| 434 |
+
- Start with a scale factor around 1.25 and denoise strength of 0.6
|
| 435 |
+
- Lower denoise values preserve more of the original video structure
|
| 436 |
+
- The second pass will increase generation time substantially - use with caution!
|
| 437 |
+
|
| 438 |
## Using LoRAs with multiple safetensors files:
|
| 439 |
If you encounter an error stating "more than one weights file", you need to specify the exact weight file name in the "LoRA Weight Name" field.
|
| 440 |
You can find this by browsing the repository on Hugging Face and looking for the safetensors files (common names include: adapter_model.safetensors, pytorch_lora_weights.safetensors).
|