Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -1,724 +1,773 @@
|
|
1 |
-
import os
|
2 |
-
|
3 |
-
os.environ['HF_HOME'] = os.path.abspath(
|
4 |
-
os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download'))
|
5 |
-
)
|
6 |
-
|
7 |
-
import gradio as gr
|
8 |
-
import torch
|
9 |
-
import traceback
|
10 |
-
import einops
|
11 |
-
import safetensors.torch as sf
|
12 |
-
import numpy as np
|
13 |
-
import math
|
14 |
-
import spaces
|
15 |
-
|
16 |
-
from PIL import Image
|
17 |
-
from diffusers import AutoencoderKLHunyuanVideo
|
18 |
-
from transformers import (
|
19 |
-
LlamaModel, CLIPTextModel,
|
20 |
-
LlamaTokenizerFast, CLIPTokenizer
|
21 |
-
)
|
22 |
-
from diffusers_helper.hunyuan import (
|
23 |
-
encode_prompt_conds, vae_decode,
|
24 |
-
vae_encode, vae_decode_fake
|
25 |
-
)
|
26 |
-
from diffusers_helper.utils import (
|
27 |
-
save_bcthw_as_mp4, crop_or_pad_yield_mask,
|
28 |
-
soft_append_bcthw, resize_and_center_crop,
|
29 |
-
state_dict_weighted_merge, state_dict_offset_merge,
|
30 |
-
generate_timestamp
|
31 |
-
)
|
32 |
-
from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
|
33 |
-
from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
|
34 |
-
from diffusers_helper.memory import (
|
35 |
-
cpu, gpu,
|
36 |
-
get_cuda_free_memory_gb,
|
37 |
-
move_model_to_device_with_memory_preservation,
|
38 |
-
offload_model_from_device_for_memory_preservation,
|
39 |
-
fake_diffusers_current_device,
|
40 |
-
DynamicSwapInstaller,
|
41 |
-
unload_complete_models,
|
42 |
-
load_model_as_complete
|
43 |
-
)
|
44 |
-
from diffusers_helper.thread_utils import AsyncStream, async_run
|
45 |
-
from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
|
46 |
-
from transformers import SiglipImageProcessor, SiglipVisionModel
|
47 |
-
from diffusers_helper.clip_vision import hf_clip_vision_encode
|
48 |
-
from diffusers_helper.bucket_tools import find_nearest_bucket
|
49 |
-
|
50 |
-
# Check GPU memory
|
51 |
-
free_mem_gb = get_cuda_free_memory_gb(gpu)
|
52 |
-
high_vram = free_mem_gb > 60
|
53 |
-
|
54 |
-
print(f'Free VRAM {free_mem_gb} GB')
|
55 |
-
print(f'High-VRAM Mode: {high_vram}')
|
56 |
-
|
57 |
-
# Load models
|
58 |
-
text_encoder = LlamaModel.from_pretrained(
|
59 |
-
"hunyuanvideo-community/HunyuanVideo",
|
60 |
-
subfolder='text_encoder',
|
61 |
-
torch_dtype=torch.float16
|
62 |
-
).cpu()
|
63 |
-
text_encoder_2 = CLIPTextModel.from_pretrained(
|
64 |
-
"hunyuanvideo-community/HunyuanVideo",
|
65 |
-
subfolder='text_encoder_2',
|
66 |
-
torch_dtype=torch.float16
|
67 |
-
).cpu()
|
68 |
-
tokenizer = LlamaTokenizerFast.from_pretrained(
|
69 |
-
"hunyuanvideo-community/HunyuanVideo",
|
70 |
-
subfolder='tokenizer'
|
71 |
-
)
|
72 |
-
tokenizer_2 = CLIPTokenizer.from_pretrained(
|
73 |
-
"hunyuanvideo-community/HunyuanVideo",
|
74 |
-
subfolder='tokenizer_2'
|
75 |
-
)
|
76 |
-
vae = AutoencoderKLHunyuanVideo.from_pretrained(
|
77 |
-
"hunyuanvideo-community/HunyuanVideo",
|
78 |
-
subfolder='vae',
|
79 |
-
torch_dtype=torch.float16
|
80 |
-
).cpu()
|
81 |
-
|
82 |
-
feature_extractor = SiglipImageProcessor.from_pretrained(
|
83 |
-
"lllyasviel/flux_redux_bfl",
|
84 |
-
subfolder='feature_extractor'
|
85 |
-
)
|
86 |
-
image_encoder = SiglipVisionModel.from_pretrained(
|
87 |
-
"lllyasviel/flux_redux_bfl",
|
88 |
-
subfolder='image_encoder',
|
89 |
-
torch_dtype=torch.float16
|
90 |
-
).cpu()
|
91 |
-
|
92 |
-
transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
|
93 |
-
'lllyasviel/FramePack_F1_I2V_HY_20250503',
|
94 |
-
torch_dtype=torch.bfloat16
|
95 |
-
).cpu()
|
96 |
-
|
97 |
-
# Evaluation mode
|
98 |
-
vae.eval()
|
99 |
-
text_encoder.eval()
|
100 |
-
text_encoder_2.eval()
|
101 |
-
image_encoder.eval()
|
102 |
-
transformer.eval()
|
103 |
-
|
104 |
-
# Slicing/Tiling for low VRAM
|
105 |
-
if not high_vram:
|
106 |
-
vae.enable_slicing()
|
107 |
-
vae.enable_tiling()
|
108 |
-
|
109 |
-
transformer.high_quality_fp32_output_for_inference = True
|
110 |
-
print('transformer.high_quality_fp32_output_for_inference = True')
|
111 |
-
|
112 |
-
# Move to correct dtype
|
113 |
-
transformer.to(dtype=torch.bfloat16)
|
114 |
-
vae.to(dtype=torch.float16)
|
115 |
-
image_encoder.to(dtype=torch.float16)
|
116 |
-
text_encoder.to(dtype=torch.float16)
|
117 |
-
text_encoder_2.to(dtype=torch.float16)
|
118 |
-
|
119 |
-
# No gradient
|
120 |
-
vae.requires_grad_(False)
|
121 |
-
text_encoder.requires_grad_(False)
|
122 |
-
text_encoder_2.requires_grad_(False)
|
123 |
-
image_encoder.requires_grad_(False)
|
124 |
-
transformer.requires_grad_(False)
|
125 |
-
|
126 |
-
# DynamicSwap if low VRAM
|
127 |
-
if not high_vram:
|
128 |
-
DynamicSwapInstaller.install_model(transformer, device=gpu)
|
129 |
-
DynamicSwapInstaller.install_model(text_encoder, device=gpu)
|
130 |
-
else:
|
131 |
-
text_encoder.to(gpu)
|
132 |
-
text_encoder_2.to(gpu)
|
133 |
-
image_encoder.to(gpu)
|
134 |
-
vae.to(gpu)
|
135 |
-
transformer.to(gpu)
|
136 |
-
|
137 |
-
stream = AsyncStream()
|
138 |
-
|
139 |
-
outputs_folder = './outputs/'
|
140 |
-
os.makedirs(outputs_folder, exist_ok=True)
|
141 |
-
|
142 |
-
examples = [
|
143 |
-
["img_examples/1.png", "The girl dances gracefully, with clear movements, full of charm."],
|
144 |
-
["img_examples/2.jpg", "The man dances flamboyantly, swinging his hips and striking bold poses with dramatic flair."],
|
145 |
-
["img_examples/3.png", "The woman dances elegantly among the blossoms, spinning slowly with flowing sleeves and graceful hand movements."]
|
146 |
-
]
|
147 |
-
|
148 |
-
# Example generation (optional)
|
149 |
-
def generate_examples(input_image, prompt):
|
150 |
-
t2v=False
|
151 |
-
n_prompt=""
|
152 |
-
seed=31337
|
153 |
-
total_second_length=60
|
154 |
-
latent_window_size=9
|
155 |
-
steps=25
|
156 |
-
cfg=1.0
|
157 |
-
gs=10.0
|
158 |
-
rs=0.0
|
159 |
-
gpu_memory_preservation=6
|
160 |
-
use_teacache=True
|
161 |
-
mp4_crf=16
|
162 |
-
|
163 |
-
global stream
|
164 |
-
|
165 |
-
if t2v:
|
166 |
-
default_height, default_width = 640, 640
|
167 |
-
input_image = np.ones((default_height, default_width, 3), dtype=np.uint8) * 255
|
168 |
-
print("No input image provided. Using a blank white image.")
|
169 |
-
|
170 |
-
yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
|
171 |
-
|
172 |
-
stream = AsyncStream()
|
173 |
-
|
174 |
-
async_run(
|
175 |
-
worker, input_image, prompt, n_prompt, seed,
|
176 |
-
total_second_length, latent_window_size, steps,
|
177 |
-
cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
|
178 |
-
)
|
179 |
-
|
180 |
-
output_filename = None
|
181 |
-
|
182 |
-
while True:
|
183 |
-
flag, data = stream.output_queue.next()
|
184 |
-
|
185 |
-
if flag == 'file':
|
186 |
-
output_filename = data
|
187 |
-
yield (
|
188 |
-
output_filename,
|
189 |
-
gr.update(),
|
190 |
-
gr.update(),
|
191 |
-
gr.update(),
|
192 |
-
gr.update(interactive=False),
|
193 |
-
gr.update(interactive=True)
|
194 |
-
)
|
195 |
-
|
196 |
-
if flag == 'progress':
|
197 |
-
preview, desc, html = data
|
198 |
-
yield (
|
199 |
-
gr.update(),
|
200 |
-
gr.update(visible=True, value=preview),
|
201 |
-
desc,
|
202 |
-
html,
|
203 |
-
gr.update(interactive=False),
|
204 |
-
gr.update(interactive=True)
|
205 |
-
)
|
206 |
-
|
207 |
-
if flag == 'end':
|
208 |
-
yield (
|
209 |
-
output_filename,
|
210 |
-
gr.update(visible=False),
|
211 |
-
gr.update(),
|
212 |
-
'',
|
213 |
-
gr.update(interactive=True),
|
214 |
-
gr.update(interactive=False)
|
215 |
-
)
|
216 |
-
break
|
217 |
-
|
218 |
-
@torch.no_grad()
|
219 |
-
def worker(
|
220 |
-
input_image, prompt, n_prompt, seed,
|
221 |
-
total_second_length, latent_window_size, steps,
|
222 |
-
cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
|
223 |
-
):
|
224 |
-
# Calculate total sections
|
225 |
-
total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
|
226 |
-
total_latent_sections = int(max(round(total_latent_sections), 1))
|
227 |
-
|
228 |
-
job_id = generate_timestamp()
|
229 |
-
|
230 |
-
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
|
231 |
-
|
232 |
-
try:
|
233 |
-
# Unload if VRAM is low
|
234 |
-
if not high_vram:
|
235 |
-
unload_complete_models(
|
236 |
-
text_encoder, text_encoder_2, image_encoder, vae, transformer
|
237 |
-
)
|
238 |
-
|
239 |
-
# Text encoding
|
240 |
-
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
|
241 |
-
|
242 |
-
if not high_vram:
|
243 |
-
fake_diffusers_current_device(text_encoder, gpu)
|
244 |
-
load_model_as_complete(text_encoder_2, target_device=gpu)
|
245 |
-
|
246 |
-
llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
|
247 |
-
|
248 |
-
if cfg == 1:
|
249 |
-
llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
|
250 |
-
else:
|
251 |
-
llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
|
252 |
-
|
253 |
-
llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
|
254 |
-
llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
|
255 |
-
|
256 |
-
# Process image
|
257 |
-
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
|
258 |
-
|
259 |
-
H, W, C = input_image.shape
|
260 |
-
height, width = find_nearest_bucket(H, W, resolution=640)
|
261 |
-
input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
|
262 |
-
|
263 |
-
Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
|
264 |
-
|
265 |
-
input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
|
266 |
-
input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
|
267 |
-
|
268 |
-
# VAE encoding
|
269 |
-
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
|
270 |
-
|
271 |
-
if not high_vram:
|
272 |
-
load_model_as_complete(vae, target_device=gpu)
|
273 |
-
start_latent = vae_encode(input_image_pt, vae)
|
274 |
-
|
275 |
-
# CLIP Vision
|
276 |
-
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
|
277 |
-
|
278 |
-
if not high_vram:
|
279 |
-
load_model_as_complete(image_encoder, target_device=gpu)
|
280 |
-
image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
|
281 |
-
image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
|
282 |
-
|
283 |
-
# Convert dtype
|
284 |
-
llama_vec = llama_vec.to(transformer.dtype)
|
285 |
-
llama_vec_n = llama_vec_n.to(transformer.dtype)
|
286 |
-
clip_l_pooler = clip_l_pooler.to(transformer.dtype)
|
287 |
-
clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
|
288 |
-
image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
|
289 |
-
|
290 |
-
# Start sampling
|
291 |
-
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
|
292 |
-
|
293 |
-
rnd = torch.Generator("cpu").manual_seed(seed)
|
294 |
-
|
295 |
-
history_latents = torch.zeros(
|
296 |
-
size=(1, 16, 16 + 2 + 1, height // 8, width // 8),
|
297 |
-
dtype=torch.float32
|
298 |
-
).cpu()
|
299 |
-
history_pixels = None
|
300 |
-
|
301 |
-
# Add start_latent
|
302 |
-
history_latents = torch.cat([history_latents, start_latent.to(history_latents)], dim=2)
|
303 |
-
total_generated_latent_frames = 1
|
304 |
-
|
305 |
-
for section_index in range(total_latent_sections):
|
306 |
-
if stream.input_queue.top() == 'end':
|
307 |
-
stream.output_queue.push(('end', None))
|
308 |
-
return
|
309 |
-
|
310 |
-
print(f'section_index = {section_index}, total_latent_sections = {total_latent_sections}')
|
311 |
-
|
312 |
-
if not high_vram:
|
313 |
-
unload_complete_models()
|
314 |
-
move_model_to_device_with_memory_preservation(
|
315 |
-
transformer, target_device=gpu,
|
316 |
-
preserved_memory_gb=gpu_memory_preservation
|
317 |
-
)
|
318 |
-
|
319 |
-
if use_teacache:
|
320 |
-
transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
|
321 |
-
else:
|
322 |
-
transformer.initialize_teacache(enable_teacache=False)
|
323 |
-
|
324 |
-
def callback(d):
|
325 |
-
preview = d['denoised']
|
326 |
-
preview = vae_decode_fake(preview)
|
327 |
-
preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
|
328 |
-
preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
|
329 |
-
|
330 |
-
if stream.input_queue.top() == 'end':
|
331 |
-
stream.output_queue.push(('end', None))
|
332 |
-
raise KeyboardInterrupt('User ends the task.')
|
333 |
-
|
334 |
-
current_step = d['i'] + 1
|
335 |
-
percentage = int(100.0 * current_step / steps)
|
336 |
-
hint = f'Sampling {current_step}/{steps}'
|
337 |
-
desc = f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}'
|
338 |
-
stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
|
339 |
-
return
|
340 |
-
|
341 |
-
indices = torch.arange(
|
342 |
-
0, sum([1, 16, 2, 1, latent_window_size])
|
343 |
-
).unsqueeze(0)
|
344 |
-
(
|
345 |
-
clean_latent_indices_start,
|
346 |
-
clean_latent_4x_indices,
|
347 |
-
clean_latent_2x_indices,
|
348 |
-
clean_latent_1x_indices,
|
349 |
-
latent_indices
|
350 |
-
) = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
|
351 |
-
|
352 |
-
clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
|
353 |
-
|
354 |
-
clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[
|
355 |
-
:, :, -sum([16, 2, 1]):, :, :
|
356 |
-
].split([16, 2, 1], dim=2)
|
357 |
-
|
358 |
-
clean_latents = torch.cat(
|
359 |
-
[start_latent.to(history_latents), clean_latents_1x],
|
360 |
-
dim=2
|
361 |
-
)
|
362 |
-
|
363 |
-
generated_latents = sample_hunyuan(
|
364 |
-
transformer=transformer,
|
365 |
-
sampler='unipc',
|
366 |
-
width=width,
|
367 |
-
height=height,
|
368 |
-
frames=latent_window_size * 4 - 3,
|
369 |
-
real_guidance_scale=cfg,
|
370 |
-
distilled_guidance_scale=gs,
|
371 |
-
guidance_rescale=rs,
|
372 |
-
num_inference_steps=steps,
|
373 |
-
generator=rnd,
|
374 |
-
prompt_embeds=llama_vec,
|
375 |
-
prompt_embeds_mask=llama_attention_mask,
|
376 |
-
prompt_poolers=clip_l_pooler,
|
377 |
-
negative_prompt_embeds=llama_vec_n,
|
378 |
-
negative_prompt_embeds_mask=llama_attention_mask_n,
|
379 |
-
negative_prompt_poolers=clip_l_pooler_n,
|
380 |
-
device=gpu,
|
381 |
-
dtype=torch.bfloat16,
|
382 |
-
image_embeddings=image_encoder_last_hidden_state,
|
383 |
-
latent_indices=latent_indices,
|
384 |
-
clean_latents=clean_latents,
|
385 |
-
clean_latent_indices=clean_latent_indices,
|
386 |
-
clean_latents_2x=clean_latents_2x,
|
387 |
-
clean_latent_2x_indices=clean_latent_2x_indices,
|
388 |
-
clean_latents_4x=clean_latents_4x,
|
389 |
-
clean_latent_4x_indices=clean_latent_4x_indices,
|
390 |
-
callback=callback,
|
391 |
-
)
|
392 |
-
|
393 |
-
total_generated_latent_frames += int(generated_latents.shape[2])
|
394 |
-
history_latents = torch.cat([history_latents, generated_latents.to(history_latents)], dim=2)
|
395 |
-
|
396 |
-
if not high_vram:
|
397 |
-
offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
|
398 |
-
load_model_as_complete(vae, target_device=gpu)
|
399 |
-
|
400 |
-
real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :]
|
401 |
-
|
402 |
-
if history_pixels is None:
|
403 |
-
history_pixels = vae_decode(real_history_latents, vae).cpu()
|
404 |
-
else:
|
405 |
-
section_latent_frames = latent_window_size * 2
|
406 |
-
overlapped_frames = latent_window_size * 4 - 3
|
407 |
-
|
408 |
-
current_pixels = vae_decode(
|
409 |
-
real_history_latents[:, :, -section_latent_frames:], vae
|
410 |
-
).cpu()
|
411 |
-
history_pixels = soft_append_bcthw(
|
412 |
-
history_pixels, current_pixels, overlapped_frames
|
413 |
-
)
|
414 |
-
|
415 |
-
if not high_vram:
|
416 |
-
unload_complete_models()
|
417 |
-
|
418 |
-
output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
|
419 |
-
|
420 |
-
save_bcthw_as_mp4(history_pixels, output_filename, fps=30)
|
421 |
-
|
422 |
-
print(f'Decoded. Latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
|
423 |
-
|
424 |
-
stream.output_queue.push(('file', output_filename))
|
425 |
-
|
426 |
-
except:
|
427 |
-
traceback.print_exc()
|
428 |
-
if not high_vram:
|
429 |
-
unload_complete_models(text_encoder, text_encoder_2, image_encoder, vae, transformer)
|
430 |
-
|
431 |
-
stream.output_queue.push(('end', None))
|
432 |
-
return
|
433 |
-
|
434 |
-
def get_duration(
|
435 |
-
input_image, prompt, t2v, n_prompt,
|
436 |
-
seed, total_second_length, latent_window_size,
|
437 |
-
steps, cfg, gs, rs, gpu_memory_preservation,
|
438 |
-
use_teacache, mp4_crf
|
439 |
-
):
|
440 |
-
return total_second_length * 60
|
441 |
-
|
442 |
-
@spaces.GPU(duration=get_duration)
|
443 |
-
def process(
|
444 |
-
input_image, prompt, t2v=False, n_prompt="", seed=31337,
|
445 |
-
total_second_length=60, latent_window_size=9, steps=25,
|
446 |
-
cfg=1.0, gs=10.0, rs=0.0, gpu_memory_preservation=6,
|
447 |
-
use_teacache=True, mp4_crf=16
|
448 |
-
):
|
449 |
-
global stream
|
450 |
-
if t2v:
|
451 |
-
default_height, default_width = 640, 640
|
452 |
-
input_image = np.ones((default_height, default_width, 3), dtype=np.uint8) * 255
|
453 |
-
print("No input image provided. Using a blank white image.")
|
454 |
-
else:
|
455 |
-
composite_rgba_uint8 = input_image["composite"]
|
456 |
-
|
457 |
-
rgb_uint8 = composite_rgba_uint8[:, :, :3]
|
458 |
-
mask_uint8 = composite_rgba_uint8[:, :, 3]
|
459 |
-
|
460 |
-
h, w = rgb_uint8.shape[:2]
|
461 |
-
background_uint8 = np.full((h, w, 3), 255, dtype=np.uint8)
|
462 |
-
|
463 |
-
alpha_normalized_float32 = mask_uint8.astype(np.float32) / 255.0
|
464 |
-
alpha_mask_float32 = np.stack([alpha_normalized_float32]*3, axis=2)
|
465 |
-
|
466 |
-
blended_image_float32 = rgb_uint8.astype(np.float32) * alpha_mask_float32 + \
|
467 |
-
background_uint8.astype(np.float32) * (1.0 - alpha_mask_float32)
|
468 |
-
|
469 |
-
input_image = np.clip(blended_image_float32, 0, 255).astype(np.uint8)
|
470 |
-
|
471 |
-
yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
|
472 |
-
|
473 |
-
stream = AsyncStream()
|
474 |
-
|
475 |
-
async_run(
|
476 |
-
worker, input_image, prompt, n_prompt, seed,
|
477 |
-
total_second_length, latent_window_size, steps,
|
478 |
-
cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
|
479 |
-
)
|
480 |
-
|
481 |
-
output_filename = None
|
482 |
-
|
483 |
-
while True:
|
484 |
-
flag, data = stream.output_queue.next()
|
485 |
-
|
486 |
-
if flag == 'file':
|
487 |
-
output_filename = data
|
488 |
-
yield (
|
489 |
-
output_filename,
|
490 |
-
gr.update(),
|
491 |
-
gr.update(),
|
492 |
-
gr.update(),
|
493 |
-
gr.update(interactive=False),
|
494 |
-
gr.update(interactive=True)
|
495 |
-
)
|
496 |
-
|
497 |
-
elif flag == 'progress':
|
498 |
-
preview, desc, html = data
|
499 |
-
yield (
|
500 |
-
gr.update(),
|
501 |
-
gr.update(visible=True, value=preview),
|
502 |
-
desc,
|
503 |
-
html,
|
504 |
-
gr.update(interactive=False),
|
505 |
-
gr.update(interactive=True)
|
506 |
-
)
|
507 |
-
|
508 |
-
elif flag == 'end':
|
509 |
-
yield (
|
510 |
-
output_filename,
|
511 |
-
gr.update(visible=False),
|
512 |
-
gr.update(),
|
513 |
-
'',
|
514 |
-
gr.update(interactive=True),
|
515 |
-
gr.update(interactive=False)
|
516 |
-
)
|
517 |
-
break
|
518 |
-
|
519 |
-
def end_process():
|
520 |
-
stream.input_queue.push('end')
|
521 |
-
|
522 |
-
|
523 |
-
quick_prompts = [
|
524 |
-
'The girl dances gracefully, with clear movements, full of charm.',
|
525 |
-
'A character doing some simple body movements.'
|
526 |
-
]
|
527 |
-
quick_prompts = [[x] for x in quick_prompts]
|
528 |
-
|
529 |
-
|
530 |
-
def make_custom_css():
|
531 |
-
base_progress_css = make_progress_bar_css()
|
532 |
-
extra_css = """
|
533 |
-
body {
|
534 |
-
background: #fafbfe !important;
|
535 |
-
font-family: "Noto Sans", sans-serif;
|
536 |
-
}
|
537 |
-
#title-container {
|
538 |
-
text-align: center;
|
539 |
-
padding: 20px 0;
|
540 |
-
|
541 |
-
|
542 |
-
|
543 |
-
|
544 |
-
|
545 |
-
|
546 |
-
|
547 |
-
|
548 |
-
|
549 |
-
|
550 |
-
|
551 |
-
|
552 |
-
|
553 |
-
|
554 |
-
|
555 |
-
|
556 |
-
|
557 |
-
|
558 |
-
|
559 |
-
|
560 |
-
|
561 |
-
|
562 |
-
|
563 |
-
|
564 |
-
|
565 |
-
|
566 |
-
|
567 |
-
|
568 |
-
|
569 |
-
|
570 |
-
|
571 |
-
|
572 |
-
|
573 |
-
|
574 |
-
|
575 |
-
|
576 |
-
|
577 |
-
|
578 |
-
|
579 |
-
|
580 |
-
|
581 |
-
|
582 |
-
|
583 |
-
|
584 |
-
|
585 |
-
margin-bottom:
|
586 |
-
|
587 |
-
|
588 |
-
|
589 |
-
|
590 |
-
|
591 |
-
|
592 |
-
|
593 |
-
|
594 |
-
|
595 |
-
|
596 |
-
|
597 |
-
|
598 |
-
|
599 |
-
|
600 |
-
|
601 |
-
|
602 |
-
|
603 |
-
|
604 |
-
|
605 |
-
|
606 |
-
|
607 |
-
|
608 |
-
|
609 |
-
|
610 |
-
|
611 |
-
|
612 |
-
|
613 |
-
|
614 |
-
|
615 |
-
|
616 |
-
|
617 |
-
|
618 |
-
|
619 |
-
|
620 |
-
|
621 |
-
|
622 |
-
|
623 |
-
|
624 |
-
|
625 |
-
|
626 |
-
|
627 |
-
|
628 |
-
|
629 |
-
|
630 |
-
|
631 |
-
|
632 |
-
|
633 |
-
|
634 |
-
|
635 |
-
|
636 |
-
|
637 |
-
|
638 |
-
|
639 |
-
|
640 |
-
|
641 |
-
|
642 |
-
|
643 |
-
|
644 |
-
|
645 |
-
|
646 |
-
|
647 |
-
|
648 |
-
|
649 |
-
|
650 |
-
|
651 |
-
|
652 |
-
|
653 |
-
|
654 |
-
|
655 |
-
|
656 |
-
|
657 |
-
|
658 |
-
|
659 |
-
|
660 |
-
|
661 |
-
|
662 |
-
|
663 |
-
|
664 |
-
|
665 |
-
|
666 |
-
|
667 |
-
|
668 |
-
|
669 |
-
|
670 |
-
|
671 |
-
|
672 |
-
|
673 |
-
|
674 |
-
|
675 |
-
|
676 |
-
|
677 |
-
|
678 |
-
|
679 |
-
|
680 |
-
|
681 |
-
|
682 |
-
|
683 |
-
|
684 |
-
|
685 |
-
|
686 |
-
|
687 |
-
|
688 |
-
|
689 |
-
|
690 |
-
|
691 |
-
|
692 |
-
|
693 |
-
|
694 |
-
|
695 |
-
|
696 |
-
|
697 |
-
|
698 |
-
)
|
699 |
-
|
700 |
-
|
701 |
-
|
702 |
-
|
703 |
-
|
704 |
-
|
705 |
-
|
706 |
-
|
707 |
-
|
708 |
-
|
709 |
-
|
710 |
-
|
711 |
-
|
712 |
-
|
713 |
-
|
714 |
-
|
715 |
-
|
716 |
-
|
717 |
-
|
718 |
-
|
719 |
-
|
720 |
-
|
721 |
-
|
722 |
-
|
723 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
724 |
block.launch(share=True)
|
|
|
1 |
+
import os
|
2 |
+
|
3 |
+
os.environ['HF_HOME'] = os.path.abspath(
|
4 |
+
os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download'))
|
5 |
+
)
|
6 |
+
|
7 |
+
import gradio as gr
|
8 |
+
import torch
|
9 |
+
import traceback
|
10 |
+
import einops
|
11 |
+
import safetensors.torch as sf
|
12 |
+
import numpy as np
|
13 |
+
import math
|
14 |
+
import spaces
|
15 |
+
|
16 |
+
from PIL import Image
|
17 |
+
from diffusers import AutoencoderKLHunyuanVideo
|
18 |
+
from transformers import (
|
19 |
+
LlamaModel, CLIPTextModel,
|
20 |
+
LlamaTokenizerFast, CLIPTokenizer
|
21 |
+
)
|
22 |
+
from diffusers_helper.hunyuan import (
|
23 |
+
encode_prompt_conds, vae_decode,
|
24 |
+
vae_encode, vae_decode_fake
|
25 |
+
)
|
26 |
+
from diffusers_helper.utils import (
|
27 |
+
save_bcthw_as_mp4, crop_or_pad_yield_mask,
|
28 |
+
soft_append_bcthw, resize_and_center_crop,
|
29 |
+
state_dict_weighted_merge, state_dict_offset_merge,
|
30 |
+
generate_timestamp
|
31 |
+
)
|
32 |
+
from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
|
33 |
+
from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
|
34 |
+
from diffusers_helper.memory import (
|
35 |
+
cpu, gpu,
|
36 |
+
get_cuda_free_memory_gb,
|
37 |
+
move_model_to_device_with_memory_preservation,
|
38 |
+
offload_model_from_device_for_memory_preservation,
|
39 |
+
fake_diffusers_current_device,
|
40 |
+
DynamicSwapInstaller,
|
41 |
+
unload_complete_models,
|
42 |
+
load_model_as_complete
|
43 |
+
)
|
44 |
+
from diffusers_helper.thread_utils import AsyncStream, async_run
|
45 |
+
from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
|
46 |
+
from transformers import SiglipImageProcessor, SiglipVisionModel
|
47 |
+
from diffusers_helper.clip_vision import hf_clip_vision_encode
|
48 |
+
from diffusers_helper.bucket_tools import find_nearest_bucket
|
49 |
+
|
50 |
+
# Check GPU memory
|
51 |
+
free_mem_gb = get_cuda_free_memory_gb(gpu)
|
52 |
+
high_vram = free_mem_gb > 60
|
53 |
+
|
54 |
+
print(f'Free VRAM {free_mem_gb} GB')
|
55 |
+
print(f'High-VRAM Mode: {high_vram}')
|
56 |
+
|
57 |
+
# Load models
|
58 |
+
text_encoder = LlamaModel.from_pretrained(
|
59 |
+
"hunyuanvideo-community/HunyuanVideo",
|
60 |
+
subfolder='text_encoder',
|
61 |
+
torch_dtype=torch.float16
|
62 |
+
).cpu()
|
63 |
+
text_encoder_2 = CLIPTextModel.from_pretrained(
|
64 |
+
"hunyuanvideo-community/HunyuanVideo",
|
65 |
+
subfolder='text_encoder_2',
|
66 |
+
torch_dtype=torch.float16
|
67 |
+
).cpu()
|
68 |
+
tokenizer = LlamaTokenizerFast.from_pretrained(
|
69 |
+
"hunyuanvideo-community/HunyuanVideo",
|
70 |
+
subfolder='tokenizer'
|
71 |
+
)
|
72 |
+
tokenizer_2 = CLIPTokenizer.from_pretrained(
|
73 |
+
"hunyuanvideo-community/HunyuanVideo",
|
74 |
+
subfolder='tokenizer_2'
|
75 |
+
)
|
76 |
+
vae = AutoencoderKLHunyuanVideo.from_pretrained(
|
77 |
+
"hunyuanvideo-community/HunyuanVideo",
|
78 |
+
subfolder='vae',
|
79 |
+
torch_dtype=torch.float16
|
80 |
+
).cpu()
|
81 |
+
|
82 |
+
feature_extractor = SiglipImageProcessor.from_pretrained(
|
83 |
+
"lllyasviel/flux_redux_bfl",
|
84 |
+
subfolder='feature_extractor'
|
85 |
+
)
|
86 |
+
image_encoder = SiglipVisionModel.from_pretrained(
|
87 |
+
"lllyasviel/flux_redux_bfl",
|
88 |
+
subfolder='image_encoder',
|
89 |
+
torch_dtype=torch.float16
|
90 |
+
).cpu()
|
91 |
+
|
92 |
+
transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
|
93 |
+
'lllyasviel/FramePack_F1_I2V_HY_20250503',
|
94 |
+
torch_dtype=torch.bfloat16
|
95 |
+
).cpu()
|
96 |
+
|
97 |
+
# Evaluation mode
|
98 |
+
vae.eval()
|
99 |
+
text_encoder.eval()
|
100 |
+
text_encoder_2.eval()
|
101 |
+
image_encoder.eval()
|
102 |
+
transformer.eval()
|
103 |
+
|
104 |
+
# Slicing/Tiling for low VRAM
|
105 |
+
if not high_vram:
|
106 |
+
vae.enable_slicing()
|
107 |
+
vae.enable_tiling()
|
108 |
+
|
109 |
+
transformer.high_quality_fp32_output_for_inference = True
|
110 |
+
print('transformer.high_quality_fp32_output_for_inference = True')
|
111 |
+
|
112 |
+
# Move to correct dtype
|
113 |
+
transformer.to(dtype=torch.bfloat16)
|
114 |
+
vae.to(dtype=torch.float16)
|
115 |
+
image_encoder.to(dtype=torch.float16)
|
116 |
+
text_encoder.to(dtype=torch.float16)
|
117 |
+
text_encoder_2.to(dtype=torch.float16)
|
118 |
+
|
119 |
+
# No gradient
|
120 |
+
vae.requires_grad_(False)
|
121 |
+
text_encoder.requires_grad_(False)
|
122 |
+
text_encoder_2.requires_grad_(False)
|
123 |
+
image_encoder.requires_grad_(False)
|
124 |
+
transformer.requires_grad_(False)
|
125 |
+
|
126 |
+
# DynamicSwap if low VRAM
|
127 |
+
if not high_vram:
|
128 |
+
DynamicSwapInstaller.install_model(transformer, device=gpu)
|
129 |
+
DynamicSwapInstaller.install_model(text_encoder, device=gpu)
|
130 |
+
else:
|
131 |
+
text_encoder.to(gpu)
|
132 |
+
text_encoder_2.to(gpu)
|
133 |
+
image_encoder.to(gpu)
|
134 |
+
vae.to(gpu)
|
135 |
+
transformer.to(gpu)
|
136 |
+
|
137 |
+
stream = AsyncStream()
|
138 |
+
|
139 |
+
outputs_folder = './outputs/'
|
140 |
+
os.makedirs(outputs_folder, exist_ok=True)
|
141 |
+
|
142 |
+
examples = [
|
143 |
+
["img_examples/1.png", "The girl dances gracefully, with clear movements, full of charm."],
|
144 |
+
["img_examples/2.jpg", "The man dances flamboyantly, swinging his hips and striking bold poses with dramatic flair."],
|
145 |
+
["img_examples/3.png", "The woman dances elegantly among the blossoms, spinning slowly with flowing sleeves and graceful hand movements."]
|
146 |
+
]
|
147 |
+
|
148 |
+
# Example generation (optional)
|
149 |
+
def generate_examples(input_image, prompt):
|
150 |
+
t2v=False
|
151 |
+
n_prompt=""
|
152 |
+
seed=31337
|
153 |
+
total_second_length=60
|
154 |
+
latent_window_size=9
|
155 |
+
steps=25
|
156 |
+
cfg=1.0
|
157 |
+
gs=10.0
|
158 |
+
rs=0.0
|
159 |
+
gpu_memory_preservation=6
|
160 |
+
use_teacache=True
|
161 |
+
mp4_crf=16
|
162 |
+
|
163 |
+
global stream
|
164 |
+
|
165 |
+
if t2v:
|
166 |
+
default_height, default_width = 640, 640
|
167 |
+
input_image = np.ones((default_height, default_width, 3), dtype=np.uint8) * 255
|
168 |
+
print("No input image provided. Using a blank white image.")
|
169 |
+
|
170 |
+
yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
|
171 |
+
|
172 |
+
stream = AsyncStream()
|
173 |
+
|
174 |
+
async_run(
|
175 |
+
worker, input_image, prompt, n_prompt, seed,
|
176 |
+
total_second_length, latent_window_size, steps,
|
177 |
+
cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
|
178 |
+
)
|
179 |
+
|
180 |
+
output_filename = None
|
181 |
+
|
182 |
+
while True:
|
183 |
+
flag, data = stream.output_queue.next()
|
184 |
+
|
185 |
+
if flag == 'file':
|
186 |
+
output_filename = data
|
187 |
+
yield (
|
188 |
+
output_filename,
|
189 |
+
gr.update(),
|
190 |
+
gr.update(),
|
191 |
+
gr.update(),
|
192 |
+
gr.update(interactive=False),
|
193 |
+
gr.update(interactive=True)
|
194 |
+
)
|
195 |
+
|
196 |
+
if flag == 'progress':
|
197 |
+
preview, desc, html = data
|
198 |
+
yield (
|
199 |
+
gr.update(),
|
200 |
+
gr.update(visible=True, value=preview),
|
201 |
+
desc,
|
202 |
+
html,
|
203 |
+
gr.update(interactive=False),
|
204 |
+
gr.update(interactive=True)
|
205 |
+
)
|
206 |
+
|
207 |
+
if flag == 'end':
|
208 |
+
yield (
|
209 |
+
output_filename,
|
210 |
+
gr.update(visible=False),
|
211 |
+
gr.update(),
|
212 |
+
'',
|
213 |
+
gr.update(interactive=True),
|
214 |
+
gr.update(interactive=False)
|
215 |
+
)
|
216 |
+
break
|
217 |
+
|
218 |
+
@torch.no_grad()
|
219 |
+
def worker(
|
220 |
+
input_image, prompt, n_prompt, seed,
|
221 |
+
total_second_length, latent_window_size, steps,
|
222 |
+
cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
|
223 |
+
):
|
224 |
+
# Calculate total sections
|
225 |
+
total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
|
226 |
+
total_latent_sections = int(max(round(total_latent_sections), 1))
|
227 |
+
|
228 |
+
job_id = generate_timestamp()
|
229 |
+
|
230 |
+
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
|
231 |
+
|
232 |
+
try:
|
233 |
+
# Unload if VRAM is low
|
234 |
+
if not high_vram:
|
235 |
+
unload_complete_models(
|
236 |
+
text_encoder, text_encoder_2, image_encoder, vae, transformer
|
237 |
+
)
|
238 |
+
|
239 |
+
# Text encoding
|
240 |
+
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
|
241 |
+
|
242 |
+
if not high_vram:
|
243 |
+
fake_diffusers_current_device(text_encoder, gpu)
|
244 |
+
load_model_as_complete(text_encoder_2, target_device=gpu)
|
245 |
+
|
246 |
+
llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
|
247 |
+
|
248 |
+
if cfg == 1:
|
249 |
+
llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
|
250 |
+
else:
|
251 |
+
llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
|
252 |
+
|
253 |
+
llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
|
254 |
+
llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
|
255 |
+
|
256 |
+
# Process image
|
257 |
+
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
|
258 |
+
|
259 |
+
H, W, C = input_image.shape
|
260 |
+
height, width = find_nearest_bucket(H, W, resolution=640)
|
261 |
+
input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
|
262 |
+
|
263 |
+
Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
|
264 |
+
|
265 |
+
input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
|
266 |
+
input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
|
267 |
+
|
268 |
+
# VAE encoding
|
269 |
+
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
|
270 |
+
|
271 |
+
if not high_vram:
|
272 |
+
load_model_as_complete(vae, target_device=gpu)
|
273 |
+
start_latent = vae_encode(input_image_pt, vae)
|
274 |
+
|
275 |
+
# CLIP Vision
|
276 |
+
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
|
277 |
+
|
278 |
+
if not high_vram:
|
279 |
+
load_model_as_complete(image_encoder, target_device=gpu)
|
280 |
+
image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
|
281 |
+
image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
|
282 |
+
|
283 |
+
# Convert dtype
|
284 |
+
llama_vec = llama_vec.to(transformer.dtype)
|
285 |
+
llama_vec_n = llama_vec_n.to(transformer.dtype)
|
286 |
+
clip_l_pooler = clip_l_pooler.to(transformer.dtype)
|
287 |
+
clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
|
288 |
+
image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
|
289 |
+
|
290 |
+
# Start sampling
|
291 |
+
stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
|
292 |
+
|
293 |
+
rnd = torch.Generator("cpu").manual_seed(seed)
|
294 |
+
|
295 |
+
history_latents = torch.zeros(
|
296 |
+
size=(1, 16, 16 + 2 + 1, height // 8, width // 8),
|
297 |
+
dtype=torch.float32
|
298 |
+
).cpu()
|
299 |
+
history_pixels = None
|
300 |
+
|
301 |
+
# Add start_latent
|
302 |
+
history_latents = torch.cat([history_latents, start_latent.to(history_latents)], dim=2)
|
303 |
+
total_generated_latent_frames = 1
|
304 |
+
|
305 |
+
for section_index in range(total_latent_sections):
|
306 |
+
if stream.input_queue.top() == 'end':
|
307 |
+
stream.output_queue.push(('end', None))
|
308 |
+
return
|
309 |
+
|
310 |
+
print(f'section_index = {section_index}, total_latent_sections = {total_latent_sections}')
|
311 |
+
|
312 |
+
if not high_vram:
|
313 |
+
unload_complete_models()
|
314 |
+
move_model_to_device_with_memory_preservation(
|
315 |
+
transformer, target_device=gpu,
|
316 |
+
preserved_memory_gb=gpu_memory_preservation
|
317 |
+
)
|
318 |
+
|
319 |
+
if use_teacache:
|
320 |
+
transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
|
321 |
+
else:
|
322 |
+
transformer.initialize_teacache(enable_teacache=False)
|
323 |
+
|
324 |
+
def callback(d):
|
325 |
+
preview = d['denoised']
|
326 |
+
preview = vae_decode_fake(preview)
|
327 |
+
preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
|
328 |
+
preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
|
329 |
+
|
330 |
+
if stream.input_queue.top() == 'end':
|
331 |
+
stream.output_queue.push(('end', None))
|
332 |
+
raise KeyboardInterrupt('User ends the task.')
|
333 |
+
|
334 |
+
current_step = d['i'] + 1
|
335 |
+
percentage = int(100.0 * current_step / steps)
|
336 |
+
hint = f'Sampling {current_step}/{steps}'
|
337 |
+
desc = f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}'
|
338 |
+
stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
|
339 |
+
return
|
340 |
+
|
341 |
+
indices = torch.arange(
|
342 |
+
0, sum([1, 16, 2, 1, latent_window_size])
|
343 |
+
).unsqueeze(0)
|
344 |
+
(
|
345 |
+
clean_latent_indices_start,
|
346 |
+
clean_latent_4x_indices,
|
347 |
+
clean_latent_2x_indices,
|
348 |
+
clean_latent_1x_indices,
|
349 |
+
latent_indices
|
350 |
+
) = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
|
351 |
+
|
352 |
+
clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
|
353 |
+
|
354 |
+
clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[
|
355 |
+
:, :, -sum([16, 2, 1]):, :, :
|
356 |
+
].split([16, 2, 1], dim=2)
|
357 |
+
|
358 |
+
clean_latents = torch.cat(
|
359 |
+
[start_latent.to(history_latents), clean_latents_1x],
|
360 |
+
dim=2
|
361 |
+
)
|
362 |
+
|
363 |
+
generated_latents = sample_hunyuan(
|
364 |
+
transformer=transformer,
|
365 |
+
sampler='unipc',
|
366 |
+
width=width,
|
367 |
+
height=height,
|
368 |
+
frames=latent_window_size * 4 - 3,
|
369 |
+
real_guidance_scale=cfg,
|
370 |
+
distilled_guidance_scale=gs,
|
371 |
+
guidance_rescale=rs,
|
372 |
+
num_inference_steps=steps,
|
373 |
+
generator=rnd,
|
374 |
+
prompt_embeds=llama_vec,
|
375 |
+
prompt_embeds_mask=llama_attention_mask,
|
376 |
+
prompt_poolers=clip_l_pooler,
|
377 |
+
negative_prompt_embeds=llama_vec_n,
|
378 |
+
negative_prompt_embeds_mask=llama_attention_mask_n,
|
379 |
+
negative_prompt_poolers=clip_l_pooler_n,
|
380 |
+
device=gpu,
|
381 |
+
dtype=torch.bfloat16,
|
382 |
+
image_embeddings=image_encoder_last_hidden_state,
|
383 |
+
latent_indices=latent_indices,
|
384 |
+
clean_latents=clean_latents,
|
385 |
+
clean_latent_indices=clean_latent_indices,
|
386 |
+
clean_latents_2x=clean_latents_2x,
|
387 |
+
clean_latent_2x_indices=clean_latent_2x_indices,
|
388 |
+
clean_latents_4x=clean_latents_4x,
|
389 |
+
clean_latent_4x_indices=clean_latent_4x_indices,
|
390 |
+
callback=callback,
|
391 |
+
)
|
392 |
+
|
393 |
+
total_generated_latent_frames += int(generated_latents.shape[2])
|
394 |
+
history_latents = torch.cat([history_latents, generated_latents.to(history_latents)], dim=2)
|
395 |
+
|
396 |
+
if not high_vram:
|
397 |
+
offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
|
398 |
+
load_model_as_complete(vae, target_device=gpu)
|
399 |
+
|
400 |
+
real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :]
|
401 |
+
|
402 |
+
if history_pixels is None:
|
403 |
+
history_pixels = vae_decode(real_history_latents, vae).cpu()
|
404 |
+
else:
|
405 |
+
section_latent_frames = latent_window_size * 2
|
406 |
+
overlapped_frames = latent_window_size * 4 - 3
|
407 |
+
|
408 |
+
current_pixels = vae_decode(
|
409 |
+
real_history_latents[:, :, -section_latent_frames:], vae
|
410 |
+
).cpu()
|
411 |
+
history_pixels = soft_append_bcthw(
|
412 |
+
history_pixels, current_pixels, overlapped_frames
|
413 |
+
)
|
414 |
+
|
415 |
+
if not high_vram:
|
416 |
+
unload_complete_models()
|
417 |
+
|
418 |
+
output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
|
419 |
+
|
420 |
+
save_bcthw_as_mp4(history_pixels, output_filename, fps=30)
|
421 |
+
|
422 |
+
print(f'Decoded. Latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
|
423 |
+
|
424 |
+
stream.output_queue.push(('file', output_filename))
|
425 |
+
|
426 |
+
except:
|
427 |
+
traceback.print_exc()
|
428 |
+
if not high_vram:
|
429 |
+
unload_complete_models(text_encoder, text_encoder_2, image_encoder, vae, transformer)
|
430 |
+
|
431 |
+
stream.output_queue.push(('end', None))
|
432 |
+
return
|
433 |
+
|
434 |
+
def get_duration(
|
435 |
+
input_image, prompt, t2v, n_prompt,
|
436 |
+
seed, total_second_length, latent_window_size,
|
437 |
+
steps, cfg, gs, rs, gpu_memory_preservation,
|
438 |
+
use_teacache, mp4_crf
|
439 |
+
):
|
440 |
+
return total_second_length * 60
|
441 |
+
|
442 |
+
@spaces.GPU(duration=get_duration)
|
443 |
+
def process(
|
444 |
+
input_image, prompt, t2v=False, n_prompt="", seed=31337,
|
445 |
+
total_second_length=60, latent_window_size=9, steps=25,
|
446 |
+
cfg=1.0, gs=10.0, rs=0.0, gpu_memory_preservation=6,
|
447 |
+
use_teacache=True, mp4_crf=16
|
448 |
+
):
|
449 |
+
global stream
|
450 |
+
if t2v:
|
451 |
+
default_height, default_width = 640, 640
|
452 |
+
input_image = np.ones((default_height, default_width, 3), dtype=np.uint8) * 255
|
453 |
+
print("No input image provided. Using a blank white image.")
|
454 |
+
else:
|
455 |
+
composite_rgba_uint8 = input_image["composite"]
|
456 |
+
|
457 |
+
rgb_uint8 = composite_rgba_uint8[:, :, :3]
|
458 |
+
mask_uint8 = composite_rgba_uint8[:, :, 3]
|
459 |
+
|
460 |
+
h, w = rgb_uint8.shape[:2]
|
461 |
+
background_uint8 = np.full((h, w, 3), 255, dtype=np.uint8)
|
462 |
+
|
463 |
+
alpha_normalized_float32 = mask_uint8.astype(np.float32) / 255.0
|
464 |
+
alpha_mask_float32 = np.stack([alpha_normalized_float32]*3, axis=2)
|
465 |
+
|
466 |
+
blended_image_float32 = rgb_uint8.astype(np.float32) * alpha_mask_float32 + \
|
467 |
+
background_uint8.astype(np.float32) * (1.0 - alpha_mask_float32)
|
468 |
+
|
469 |
+
input_image = np.clip(blended_image_float32, 0, 255).astype(np.uint8)
|
470 |
+
|
471 |
+
yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
|
472 |
+
|
473 |
+
stream = AsyncStream()
|
474 |
+
|
475 |
+
async_run(
|
476 |
+
worker, input_image, prompt, n_prompt, seed,
|
477 |
+
total_second_length, latent_window_size, steps,
|
478 |
+
cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
|
479 |
+
)
|
480 |
+
|
481 |
+
output_filename = None
|
482 |
+
|
483 |
+
while True:
|
484 |
+
flag, data = stream.output_queue.next()
|
485 |
+
|
486 |
+
if flag == 'file':
|
487 |
+
output_filename = data
|
488 |
+
yield (
|
489 |
+
output_filename,
|
490 |
+
gr.update(),
|
491 |
+
gr.update(),
|
492 |
+
gr.update(),
|
493 |
+
gr.update(interactive=False),
|
494 |
+
gr.update(interactive=True)
|
495 |
+
)
|
496 |
+
|
497 |
+
elif flag == 'progress':
|
498 |
+
preview, desc, html = data
|
499 |
+
yield (
|
500 |
+
gr.update(),
|
501 |
+
gr.update(visible=True, value=preview),
|
502 |
+
desc,
|
503 |
+
html,
|
504 |
+
gr.update(interactive=False),
|
505 |
+
gr.update(interactive=True)
|
506 |
+
)
|
507 |
+
|
508 |
+
elif flag == 'end':
|
509 |
+
yield (
|
510 |
+
output_filename,
|
511 |
+
gr.update(visible=False),
|
512 |
+
gr.update(),
|
513 |
+
'',
|
514 |
+
gr.update(interactive=True),
|
515 |
+
gr.update(interactive=False)
|
516 |
+
)
|
517 |
+
break
|
518 |
+
|
519 |
+
def end_process():
|
520 |
+
stream.input_queue.push('end')
|
521 |
+
|
522 |
+
|
523 |
+
quick_prompts = [
|
524 |
+
'The girl dances gracefully, with clear movements, full of charm.',
|
525 |
+
'A character doing some simple body movements.'
|
526 |
+
]
|
527 |
+
quick_prompts = [[x] for x in quick_prompts]
|
528 |
+
|
529 |
+
|
530 |
+
def make_custom_css():
|
531 |
+
base_progress_css = make_progress_bar_css()
|
532 |
+
extra_css = """
|
533 |
+
body {
|
534 |
+
background: #fafbfe !important;
|
535 |
+
font-family: "Noto Sans", sans-serif;
|
536 |
+
}
|
537 |
+
#title-container {
|
538 |
+
text-align: center;
|
539 |
+
padding: 20px 0;
|
540 |
+
margin-bottom: 30px;
|
541 |
+
background: linear-gradient(135deg, #4b9ffa 0%, #2d7eeb 100%);
|
542 |
+
border-radius: 15px;
|
543 |
+
box-shadow: 0 4px 15px rgba(0,0,0,0.1);
|
544 |
+
}
|
545 |
+
#title-container h1 {
|
546 |
+
color: white;
|
547 |
+
font-size: 2.5rem;
|
548 |
+
margin: 0;
|
549 |
+
font-weight: 800;
|
550 |
+
text-shadow: 1px 2px 2px rgba(0,0,0,0.2);
|
551 |
+
}
|
552 |
+
.container {
|
553 |
+
display: flex;
|
554 |
+
gap: 20px;
|
555 |
+
}
|
556 |
+
.settings-panel {
|
557 |
+
flex: 0 0 350px;
|
558 |
+
background: #ffffff;
|
559 |
+
padding: 20px;
|
560 |
+
border-radius: 15px;
|
561 |
+
box-shadow: 0 2px 8px rgba(0,0,0,0.1);
|
562 |
+
}
|
563 |
+
.settings-panel h3 {
|
564 |
+
color: #2d7eeb;
|
565 |
+
margin-bottom: 20px;
|
566 |
+
font-size: 1.3rem;
|
567 |
+
border-bottom: 2px solid #4b9ffa;
|
568 |
+
padding-bottom: 10px;
|
569 |
+
}
|
570 |
+
.main-panel {
|
571 |
+
flex: 1;
|
572 |
+
background: #ffffff;
|
573 |
+
padding: 20px;
|
574 |
+
border-radius: 15px;
|
575 |
+
box-shadow: 0 2px 8px rgba(0,0,0,0.1);
|
576 |
+
}
|
577 |
+
.gr-form {
|
578 |
+
border: none !important;
|
579 |
+
background: transparent !important;
|
580 |
+
}
|
581 |
+
.gr-box {
|
582 |
+
border: 1px solid #e0e0f0 !important;
|
583 |
+
background: #f8f9fe !important;
|
584 |
+
border-radius: 10px !important;
|
585 |
+
margin-bottom: 15px !important;
|
586 |
+
transition: all 0.3s ease;
|
587 |
+
}
|
588 |
+
.gr-box:hover {
|
589 |
+
border-color: #4b9ffa !important;
|
590 |
+
box-shadow: 0 2px 8px rgba(75, 159, 250, 0.1) !important;
|
591 |
+
}
|
592 |
+
.gr-input, .gr-button {
|
593 |
+
border-radius: 8px !important;
|
594 |
+
transition: all 0.3s ease !important;
|
595 |
+
}
|
596 |
+
.gr-button {
|
597 |
+
min-height: 45px !important;
|
598 |
+
font-weight: 600 !important;
|
599 |
+
text-transform: uppercase !important;
|
600 |
+
letter-spacing: 0.5px !important;
|
601 |
+
}
|
602 |
+
.gr-button:hover {
|
603 |
+
transform: translateY(-1px) !important;
|
604 |
+
}
|
605 |
+
.gr-button.primary-btn {
|
606 |
+
background: #4b9ffa !important;
|
607 |
+
color: white !important;
|
608 |
+
border: none !important;
|
609 |
+
}
|
610 |
+
.gr-button.secondary-btn {
|
611 |
+
background: #ff4d4d !important;
|
612 |
+
color: white !important;
|
613 |
+
border: none !important;
|
614 |
+
}
|
615 |
+
.progress-container {
|
616 |
+
margin-top: 20px;
|
617 |
+
padding: 15px;
|
618 |
+
background: #f8f9fe;
|
619 |
+
border-radius: 10px;
|
620 |
+
}
|
621 |
+
"""
|
622 |
+
return base_progress_css + extra_css
|
623 |
+
|
624 |
+
css = make_custom_css()
|
625 |
+
|
626 |
+
block = gr.Blocks(css=css).queue()
|
627 |
+
with block:
|
628 |
+
with gr.Group(elem_id="title-container"):
|
629 |
+
gr.Markdown("<h1>FramePack I2V</h1>")
|
630 |
+
gr.Markdown(
|
631 |
+
"""Generate amazing animations from a single image using AI.
|
632 |
+
Just upload an image, write a prompt, and watch the magic happen!"""
|
633 |
+
)
|
634 |
+
|
635 |
+
with gr.Row(elem_classes="container"):
|
636 |
+
with gr.Column(elem_classes="settings-panel"):
|
637 |
+
gr.Markdown("### Settings")
|
638 |
+
|
639 |
+
# Basic Settings
|
640 |
+
with gr.Group():
|
641 |
+
input_image = gr.Image(
|
642 |
+
label="Upload Image",
|
643 |
+
type="numpy",
|
644 |
+
height=320
|
645 |
+
)
|
646 |
+
prompt = gr.Textbox(
|
647 |
+
label="Describe the animation you want",
|
648 |
+
placeholder="E.g., The character dances gracefully with flowing movements...",
|
649 |
+
lines=3
|
650 |
+
)
|
651 |
+
total_second_length = gr.Slider(
|
652 |
+
label="Video Length (Seconds)",
|
653 |
+
minimum=1,
|
654 |
+
maximum=60,
|
655 |
+
value=2,
|
656 |
+
step=0.1
|
657 |
+
)
|
658 |
+
|
659 |
+
# Advanced Settings
|
660 |
+
with gr.Group():
|
661 |
+
steps = gr.Slider(
|
662 |
+
label="Generation Steps",
|
663 |
+
minimum=1,
|
664 |
+
maximum=100,
|
665 |
+
value=25,
|
666 |
+
step=1,
|
667 |
+
info='Higher values = better quality but slower'
|
668 |
+
)
|
669 |
+
gs = gr.Slider(
|
670 |
+
label="Animation Strength",
|
671 |
+
minimum=1.0,
|
672 |
+
maximum=32.0,
|
673 |
+
value=10.0,
|
674 |
+
step=0.1,
|
675 |
+
info='Controls how closely the animation follows the prompt'
|
676 |
+
)
|
677 |
+
use_teacache = gr.Checkbox(
|
678 |
+
label='Fast Mode',
|
679 |
+
value=True,
|
680 |
+
info='Faster generation but may affect quality of fine details'
|
681 |
+
)
|
682 |
+
gpu_memory_preservation = gr.Slider(
|
683 |
+
label="VRAM Usage",
|
684 |
+
minimum=6,
|
685 |
+
maximum=128,
|
686 |
+
value=6,
|
687 |
+
step=0.1,
|
688 |
+
info="Increase if you get out of memory errors"
|
689 |
+
)
|
690 |
+
seed = gr.Number(
|
691 |
+
label="Seed",
|
692 |
+
value=31337,
|
693 |
+
precision=0,
|
694 |
+
info="Change for different results"
|
695 |
+
)
|
696 |
+
|
697 |
+
# Hidden settings
|
698 |
+
n_prompt = gr.Textbox(visible=False, value="")
|
699 |
+
latent_window_size = gr.Slider(visible=False, value=9)
|
700 |
+
cfg = gr.Slider(visible=False, value=1.0)
|
701 |
+
rs = gr.Slider(visible=False, value=0.0)
|
702 |
+
|
703 |
+
with gr.Row():
|
704 |
+
start_button = gr.Button(
|
705 |
+
value="▶️ Generate Animation",
|
706 |
+
elem_classes=["primary-btn"]
|
707 |
+
)
|
708 |
+
stop_button = gr.Button(
|
709 |
+
value="⏹️ Stop",
|
710 |
+
elem_classes=["secondary-btn"],
|
711 |
+
interactive=False
|
712 |
+
)
|
713 |
+
|
714 |
+
with gr.Column(elem_classes="main-panel"):
|
715 |
+
preview_image = gr.Image(
|
716 |
+
label="Generation Preview",
|
717 |
+
height=200,
|
718 |
+
visible=False
|
719 |
+
)
|
720 |
+
result_video = gr.Video(
|
721 |
+
label="Generated Animation",
|
722 |
+
autoplay=True,
|
723 |
+
show_share_button=True,
|
724 |
+
height=512,
|
725 |
+
loop=True
|
726 |
+
)
|
727 |
+
with gr.Group(elem_classes="progress-container"):
|
728 |
+
progress_desc = gr.Markdown(
|
729 |
+
elem_classes='no-generating-animation'
|
730 |
+
)
|
731 |
+
progress_bar = gr.HTML(
|
732 |
+
elem_classes='no-generating-animation'
|
733 |
+
)
|
734 |
+
|
735 |
+
# Quick Prompts Section
|
736 |
+
with gr.Group():
|
737 |
+
gr.Markdown("### 💡 Quick Prompt Ideas")
|
738 |
+
example_quick_prompts = gr.Dataset(
|
739 |
+
samples=quick_prompts,
|
740 |
+
label='Click any prompt to try it',
|
741 |
+
samples_per_page=5,
|
742 |
+
components=[prompt]
|
743 |
+
)
|
744 |
+
|
745 |
+
# Setup callbacks
|
746 |
+
ips = [
|
747 |
+
input_image, prompt, n_prompt, seed,
|
748 |
+
total_second_length, latent_window_size,
|
749 |
+
steps, cfg, gs, rs, gpu_memory_preservation,
|
750 |
+
use_teacache, mp4_crf
|
751 |
+
]
|
752 |
+
|
753 |
+
start_button.click(
|
754 |
+
fn=process,
|
755 |
+
inputs=ips,
|
756 |
+
outputs=[
|
757 |
+
result_video, preview_image,
|
758 |
+
progress_desc, progress_bar,
|
759 |
+
start_button, stop_button
|
760 |
+
]
|
761 |
+
)
|
762 |
+
|
763 |
+
stop_button.click(fn=end_process)
|
764 |
+
|
765 |
+
example_quick_prompts.click(
|
766 |
+
fn=lambda x: x[0],
|
767 |
+
inputs=[example_quick_prompts],
|
768 |
+
outputs=prompt,
|
769 |
+
show_progress=False,
|
770 |
+
queue=False
|
771 |
+
)
|
772 |
+
|
773 |
block.launch(share=True)
|