LPX55 commited on
Commit
3ad36df
·
verified ·
1 Parent(s): c45457d

Create app_v2v.py

Browse files
Files changed (1) hide show
  1. app_v2v.py +683 -0
app_v2v.py ADDED
@@ -0,0 +1,683 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers_helper.hf_login import login
2
+
3
+ import os
4
+
5
+ os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
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 argparse
14
+ import math
15
+ # 20250506 pftq: Added for video input loading
16
+ import decord
17
+ # 20250506 pftq: Added for progress bars in video_encode
18
+ from tqdm import tqdm
19
+ # 20250506 pftq: Normalize file paths for Windows compatibility
20
+ import pathlib
21
+ # 20250506 pftq: for easier to read timestamp
22
+ from datetime import datetime
23
+ # 20250508 pftq: for saving prompt to mp4 comments metadata
24
+ import imageio_ffmpeg
25
+ import tempfile
26
+ import shutil
27
+ import subprocess
28
+
29
+ from PIL import Image
30
+ from diffusers import AutoencoderKLHunyuanVideo
31
+ from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
32
+ from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake
33
+ from diffusers_helper.utils import save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw, resize_and_center_crop, state_dict_weighted_merge, state_dict_offset_merge, generate_timestamp
34
+ from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
35
+ from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
36
+ from diffusers_helper.memory import cpu, gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation, offload_model_from_device_for_memory_preservation, fake_diffusers_current_device, DynamicSwapInstaller, unload_complete_models, load_model_as_complete
37
+ from diffusers_helper.thread_utils import AsyncStream, async_run
38
+ from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
39
+ from transformers import SiglipImageProcessor, SiglipVisionModel
40
+ from diffusers_helper.clip_vision import hf_clip_vision_encode
41
+ from diffusers_helper.bucket_tools import find_nearest_bucket
42
+
43
+ parser = argparse.ArgumentParser()
44
+ parser.add_argument('--share', action='store_true')
45
+ parser.add_argument("--server", type=str, default='0.0.0.0')
46
+ parser.add_argument("--port", type=int, required=False)
47
+ parser.add_argument("--inbrowser", action='store_true')
48
+ args = parser.parse_args()
49
+
50
+ print(args)
51
+
52
+ free_mem_gb = get_cuda_free_memory_gb(gpu)
53
+ high_vram = free_mem_gb > 60
54
+
55
+ print(f'Free VRAM {free_mem_gb} GB')
56
+ print(f'High-VRAM Mode: {high_vram}')
57
+
58
+ text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
59
+ text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
60
+ tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
61
+ tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
62
+ vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
63
+
64
+ feature_extractor = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='feature_extractor')
65
+ image_encoder = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder='image_encoder', torch_dtype=torch.float16).cpu()
66
+
67
+ transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePack_F1_I2V_HY_20250503', torch_dtype=torch.bfloat16).cpu()
68
+
69
+ vae.eval()
70
+ text_encoder.eval()
71
+ text_encoder_2.eval()
72
+ image_encoder.eval()
73
+ transformer.eval()
74
+
75
+ if not high_vram:
76
+ vae.enable_slicing()
77
+ vae.enable_tiling()
78
+
79
+ transformer.high_quality_fp32_output_for_inference = True
80
+ print('transformer.high_quality_fp32_output_for_inference = True')
81
+
82
+ transformer.to(dtype=torch.bfloat16)
83
+ vae.to(dtype=torch.float16)
84
+ image_encoder.to(dtype=torch.float16)
85
+ text_encoder.to(dtype=torch.float16)
86
+ text_encoder_2.to(dtype=torch.float16)
87
+
88
+ vae.requires_grad_(False)
89
+ text_encoder.requires_grad_(False)
90
+ text_encoder_2.requires_grad_(False)
91
+ image_encoder.requires_grad_(False)
92
+ transformer.requires_grad_(False)
93
+
94
+ if not high_vram:
95
+ # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
96
+ DynamicSwapInstaller.install_model(transformer, device=gpu)
97
+ DynamicSwapInstaller.install_model(text_encoder, device=gpu)
98
+ else:
99
+ text_encoder.to(gpu)
100
+ text_encoder_2.to(gpu)
101
+ image_encoder.to(gpu)
102
+ vae.to(gpu)
103
+ transformer.to(gpu)
104
+
105
+ stream = AsyncStream()
106
+
107
+ outputs_folder = './outputs/'
108
+ os.makedirs(outputs_folder, exist_ok=True)
109
+
110
+ # 20250506 pftq: Added function to encode input video frames into latents
111
+ @torch.no_grad()
112
+ def video_encode(video_path, resolution, no_resize, vae, vae_batch_size=16, device="cuda", width=None, height=None):
113
+ """
114
+ Encode a video into latent representations using the VAE.
115
+
116
+ Args:
117
+ video_path: Path to the input video file.
118
+ vae: AutoencoderKLHunyuanVideo model.
119
+ height, width: Target resolution for resizing frames.
120
+ vae_batch_size: Number of frames to process per batch.
121
+ device: Device for computation (e.g., "cuda").
122
+
123
+ Returns:
124
+ start_latent: Latent of the first frame (for compatibility with original code).
125
+ input_image_np: First frame as numpy array (for CLIP vision encoding).
126
+ history_latents: Latents of all frames (shape: [1, channels, frames, height//8, width//8]).
127
+ fps: Frames per second of the input video.
128
+ """
129
+ # 20250506 pftq: Normalize video path for Windows compatibility
130
+ video_path = str(pathlib.Path(video_path).resolve())
131
+ print(f"Processing video: {video_path}")
132
+
133
+ # 20250506 pftq: Check CUDA availability and fallback to CPU if needed
134
+ if device == "cuda" and not torch.cuda.is_available():
135
+ print("CUDA is not available, falling back to CPU")
136
+ device = "cpu"
137
+
138
+ try:
139
+ # 20250506 pftq: Load video and get FPS
140
+ print("Initializing VideoReader...")
141
+ vr = decord.VideoReader(video_path)
142
+ fps = vr.get_avg_fps() # Get input video FPS
143
+ num_real_frames = len(vr)
144
+ print(f"Video loaded: {num_real_frames} frames, FPS: {fps}")
145
+
146
+ # Truncate to nearest latent size (multiple of 4)
147
+ latent_size_factor = 4
148
+ num_frames = (num_real_frames // latent_size_factor) * latent_size_factor
149
+ if num_frames != num_real_frames:
150
+ print(f"Truncating video from {num_real_frames} to {num_frames} frames for latent size compatibility")
151
+ num_real_frames = num_frames
152
+
153
+ # 20250506 pftq: Read frames
154
+ print("Reading video frames...")
155
+ frames = vr.get_batch(range(num_real_frames)).asnumpy() # Shape: (num_real_frames, height, width, channels)
156
+ print(f"Frames read: {frames.shape}")
157
+
158
+ # 20250506 pftq: Get native video resolution
159
+ native_height, native_width = frames.shape[1], frames.shape[2]
160
+ print(f"Native video resolution: {native_width}x{native_height}")
161
+
162
+ # 20250506 pftq: Use native resolution if height/width not specified, otherwise use provided values
163
+ target_height = native_height if height is None else height
164
+ target_width = native_width if width is None else width
165
+
166
+ # 20250506 pftq: Adjust to nearest bucket for model compatibility
167
+ if not no_resize:
168
+ target_height, target_width = find_nearest_bucket(target_height, target_width, resolution=resolution)
169
+ print(f"Adjusted resolution: {target_width}x{target_height}")
170
+ else:
171
+ print(f"Using native resolution without resizing: {target_width}x{target_height}")
172
+
173
+ # 20250506 pftq: Preprocess frames to match original image processing
174
+ processed_frames = []
175
+ for i, frame in enumerate(frames):
176
+ #print(f"Preprocessing frame {i+1}/{num_frames}")
177
+ frame_np = resize_and_center_crop(frame, target_width=target_width, target_height=target_height)
178
+ processed_frames.append(frame_np)
179
+ processed_frames = np.stack(processed_frames) # Shape: (num_real_frames, height, width, channels)
180
+ print(f"Frames preprocessed: {processed_frames.shape}")
181
+
182
+ # 20250506 pftq: Save first frame for CLIP vision encoding
183
+ input_image_np = processed_frames[0]
184
+
185
+ # 20250506 pftq: Convert to tensor and normalize to [-1, 1]
186
+ print("Converting frames to tensor...")
187
+ frames_pt = torch.from_numpy(processed_frames).float() / 127.5 - 1
188
+ frames_pt = frames_pt.permute(0, 3, 1, 2) # Shape: (num_real_frames, channels, height, width)
189
+ frames_pt = frames_pt.unsqueeze(0) # Shape: (1, num_real_frames, channels, height, width)
190
+ frames_pt = frames_pt.permute(0, 2, 1, 3, 4) # Shape: (1, channels, num_real_frames, height, width)
191
+ print(f"Tensor shape: {frames_pt.shape}")
192
+
193
+ # 20250507 pftq: Save pixel frames for use in worker
194
+ input_video_pixels = frames_pt.cpu()
195
+
196
+ # 20250506 pftq: Move to device
197
+ print(f"Moving tensor to device: {device}")
198
+ frames_pt = frames_pt.to(device)
199
+ print("Tensor moved to device")
200
+
201
+ # 20250506 pftq: Move VAE to device
202
+ print(f"Moving VAE to device: {device}")
203
+ vae.to(device)
204
+ print("VAE moved to device")
205
+
206
+ # 20250506 pftq: Encode frames in batches
207
+ print(f"Encoding input video frames in VAE batch size {vae_batch_size} (reduce if memory issues here or if forcing video resolution)")
208
+ latents = []
209
+ vae.eval()
210
+ with torch.no_grad():
211
+ for i in tqdm(range(0, frames_pt.shape[2], vae_batch_size), desc="Encoding video frames", mininterval=0.1):
212
+ #print(f"Encoding batch {i//vae_batch_size + 1}: frames {i} to {min(i + vae_batch_size, frames_pt.shape[2])}")
213
+ batch = frames_pt[:, :, i:i + vae_batch_size] # Shape: (1, channels, batch_size, height, width)
214
+ try:
215
+ # 20250506 pftq: Log GPU memory before encoding
216
+ if device == "cuda":
217
+ free_mem = torch.cuda.memory_allocated() / 1024**3
218
+ #print(f"GPU memory before encoding: {free_mem:.2f} GB")
219
+ batch_latent = vae_encode(batch, vae)
220
+ # 20250506 pftq: Synchronize CUDA to catch issues
221
+ if device == "cuda":
222
+ torch.cuda.synchronize()
223
+ #print(f"GPU memory after encoding: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
224
+ latents.append(batch_latent)
225
+ #print(f"Batch encoded, latent shape: {batch_latent.shape}")
226
+ except RuntimeError as e:
227
+ print(f"Error during VAE encoding: {str(e)}")
228
+ if device == "cuda" and "out of memory" in str(e).lower():
229
+ print("CUDA out of memory, try reducing vae_batch_size or using CPU")
230
+ raise
231
+
232
+ # 20250506 pftq: Concatenate latents
233
+ print("Concatenating latents...")
234
+ history_latents = torch.cat(latents, dim=2) # Shape: (1, channels, frames, height//8, width//8)
235
+ print(f"History latents shape: {history_latents.shape}")
236
+
237
+ # 20250506 pftq: Get first frame's latent
238
+ start_latent = history_latents[:, :, :1] # Shape: (1, channels, 1, height//8, width//8)
239
+ print(f"Start latent shape: {start_latent.shape}")
240
+
241
+ # 20250506 pftq: Move VAE back to CPU to free GPU memory
242
+ if device == "cuda":
243
+ vae.to(cpu)
244
+ torch.cuda.empty_cache()
245
+ print("VAE moved back to CPU, CUDA cache cleared")
246
+
247
+ return start_latent, input_image_np, history_latents, fps, target_height, target_width, input_video_pixels
248
+
249
+ except Exception as e:
250
+ print(f"Error in video_encode: {str(e)}")
251
+ raise
252
+
253
+ # 20250508 pftq: for saving prompt to mp4 metadata comments
254
+ def set_mp4_comments_imageio_ffmpeg(input_file, comments):
255
+ try:
256
+ # Get the path to the bundled FFmpeg binary from imageio-ffmpeg
257
+ ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
258
+
259
+ # Check if input file exists
260
+ if not os.path.exists(input_file):
261
+ print(f"Error: Input file {input_file} does not exist")
262
+ return False
263
+
264
+ # Create a temporary file path
265
+ temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
266
+
267
+ # FFmpeg command using the bundled binary
268
+ command = [
269
+ ffmpeg_path, # Use imageio-ffmpeg's FFmpeg
270
+ '-i', input_file, # input file
271
+ '-metadata', f'comment={comments}', # set comment metadata
272
+ '-c:v', 'copy', # copy video stream without re-encoding
273
+ '-c:a', 'copy', # copy audio stream without re-encoding
274
+ '-y', # overwrite output file if it exists
275
+ temp_file # temporary output file
276
+ ]
277
+
278
+ # Run the FFmpeg command
279
+ result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
280
+
281
+ if result.returncode == 0:
282
+ # Replace the original file with the modified one
283
+ shutil.move(temp_file, input_file)
284
+ print(f"Successfully added comments to {input_file}")
285
+ return True
286
+ else:
287
+ # Clean up temp file if FFmpeg fails
288
+ if os.path.exists(temp_file):
289
+ os.remove(temp_file)
290
+ print(f"Error: FFmpeg failed with message:\n{result.stderr}")
291
+ return False
292
+
293
+ except Exception as e:
294
+ # Clean up temp file in case of other errors
295
+ if 'temp_file' in locals() and os.path.exists(temp_file):
296
+ os.remove(temp_file)
297
+ print(f"Error saving prompt to video metadata, ffmpeg may be required: "+str(e))
298
+ return False
299
+
300
+ # 20250506 pftq: Modified worker to accept video input and clean frame count
301
+ @torch.no_grad()
302
+ def worker(input_video, prompt, n_prompt, seed, batch, resolution, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, no_resize, mp4_crf, num_clean_frames, vae_batch):
303
+
304
+ stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
305
+
306
+ try:
307
+ # Clean GPU
308
+ if not high_vram:
309
+ unload_complete_models(
310
+ text_encoder, text_encoder_2, image_encoder, vae, transformer
311
+ )
312
+
313
+ # Text encoding
314
+ stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
315
+
316
+ if not high_vram:
317
+ fake_diffusers_current_device(text_encoder, gpu) # since we only encode one text - that is one model move and one encode, offload is same time consumption since it is also one load and one encode.
318
+ load_model_as_complete(text_encoder_2, target_device=gpu)
319
+
320
+ llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
321
+
322
+ if cfg == 1:
323
+ llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
324
+ else:
325
+ llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
326
+
327
+ llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
328
+ llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
329
+
330
+ # 20250506 pftq: Processing input video instead of image
331
+ stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Video processing ...'))))
332
+
333
+ # 20250506 pftq: Encode video
334
+ #H, W = 640, 640 # Default resolution, will be adjusted
335
+ #height, width = find_nearest_bucket(H, W, resolution=640)
336
+ #start_latent, input_image_np, history_latents, fps = video_encode(input_video, vae, height, width, vae_batch_size=16, device=gpu)
337
+ start_latent, input_image_np, video_latents, fps, height, width, input_video_pixels = video_encode(input_video, resolution, no_resize, vae, vae_batch_size=vae_batch, device=gpu)
338
+
339
+ #Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
340
+
341
+ # CLIP Vision
342
+ stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
343
+
344
+ if not high_vram:
345
+ load_model_as_complete(image_encoder, target_device=gpu)
346
+
347
+ image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
348
+ image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
349
+
350
+ # Dtype
351
+ llama_vec = llama_vec.to(transformer.dtype)
352
+ llama_vec_n = llama_vec_n.to(transformer.dtype)
353
+ clip_l_pooler = clip_l_pooler.to(transformer.dtype)
354
+ clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
355
+ image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
356
+
357
+ total_latent_sections = (total_second_length * fps) / (latent_window_size * 4)
358
+ total_latent_sections = int(max(round(total_latent_sections), 1))
359
+
360
+ for idx in range(batch):
361
+ if idx>0:
362
+ seed = seed + 1
363
+
364
+ if batch > 1:
365
+ print(f"Beginning video {idx+1} of {batch} with seed {seed} ")
366
+
367
+ #job_id = generate_timestamp()
368
+ job_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")+f"_framepackf1-videoinput_{width}-{total_second_length}sec_seed-{seed}_steps-{steps}_distilled-{gs}_cfg-{cfg}" # 20250506 pftq: easier to read timestamp and filename
369
+
370
+ # Sampling
371
+ stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
372
+
373
+ rnd = torch.Generator("cpu").manual_seed(seed)
374
+
375
+ # 20250506 pftq: Initialize history_latents with video latents
376
+ history_latents = video_latents.cpu()
377
+ total_generated_latent_frames = history_latents.shape[2]
378
+ # 20250506 pftq: Initialize history_pixels to fix UnboundLocalError
379
+ history_pixels = None
380
+ previous_video = None
381
+
382
+ # 20250507 pftq: hot fix for initial video being corrupted by vae encoding, issue with ghosting because of slight differences
383
+ #history_pixels = input_video_pixels
384
+ #save_bcthw_as_mp4(vae_decode(video_latents, vae).cpu(), os.path.join(outputs_folder, f'{job_id}_input_video.mp4'), fps=fps, crf=mp4_crf) # 20250507 pftq: test fast movement corrupted by vae encoding if vae batch size too low
385
+
386
+ for section_index in range(total_latent_sections):
387
+ if stream.input_queue.top() == 'end':
388
+ stream.output_queue.push(('end', None))
389
+ return
390
+
391
+ print(f'section_index = {section_index}, total_latent_sections = {total_latent_sections}')
392
+
393
+ if not high_vram:
394
+ unload_complete_models()
395
+ move_model_to_device_with_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
396
+
397
+ if use_teacache:
398
+ transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
399
+ else:
400
+ transformer.initialize_teacache(enable_teacache=False)
401
+
402
+ def callback(d):
403
+ preview = d['denoised']
404
+ preview = vae_decode_fake(preview)
405
+
406
+ preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
407
+ preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
408
+
409
+ if stream.input_queue.top() == 'end':
410
+ stream.output_queue.push(('end', None))
411
+ raise KeyboardInterrupt('User ends the task.')
412
+
413
+ current_step = d['i'] + 1
414
+ percentage = int(100.0 * current_step / steps)
415
+ hint = f'Sampling {current_step}/{steps}'
416
+ desc = f'Total frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, Video length: {max(0, (total_generated_latent_frames * 4 - 3) / fps) :.2f} seconds (FPS-{fps}), Seed: {seed}, Video {idx+1} of {batch}. The video is generating part {section_index+1} of {total_latent_sections}...'
417
+ stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
418
+ return
419
+
420
+ # 20250506 pftq: Use user-specified number of context frames, matching original allocation for num_clean_frames=2
421
+ available_frames = history_latents.shape[2] # Number of latent frames
422
+ max_pixel_frames = min(latent_window_size * 4 - 3, available_frames * 4) # Cap at available pixel frames
423
+ adjusted_latent_frames = max(1, (max_pixel_frames + 3) // 4) # Convert back to latent frames
424
+ # Adjust num_clean_frames to match original behavior: num_clean_frames=2 means 1 frame for clean_latents_1x
425
+ effective_clean_frames = max(0, num_clean_frames - 1) if num_clean_frames > 1 else 0
426
+ effective_clean_frames = min(effective_clean_frames, available_frames - 2) if available_frames > 2 else 0 # 20250507 pftq: changed 1 to 2 for edge case for <=1 sec videos
427
+ num_2x_frames = min(2, max(1, available_frames - effective_clean_frames - 1)) if available_frames > effective_clean_frames + 1 else 0 # 20250507 pftq: subtracted 1 for edge case for <=1 sec videos
428
+ num_4x_frames = min(16, max(1, available_frames - effective_clean_frames - num_2x_frames)) if available_frames > effective_clean_frames + num_2x_frames else 0 # 20250507 pftq: Edge case for <=1 sec
429
+
430
+ total_context_frames = num_4x_frames + num_2x_frames + effective_clean_frames
431
+ total_context_frames = min(total_context_frames, available_frames) # 20250507 pftq: Edge case for <=1 sec videos
432
+
433
+ indices = torch.arange(0, sum([1, num_4x_frames, num_2x_frames, effective_clean_frames, adjusted_latent_frames])).unsqueeze(0) # 20250507 pftq: latent_window_size to adjusted_latent_frames for edge case for <=1 sec videos
434
+ clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split(
435
+ [1, num_4x_frames, num_2x_frames, effective_clean_frames, adjusted_latent_frames], dim=1 # 20250507 pftq: latent_window_size to adjusted_latent_frames for edge case for <=1 sec videos
436
+ )
437
+ clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
438
+
439
+ # 20250506 pftq: Split history_latents dynamically based on available frames
440
+ fallback_frame_count = 2 # 20250507 pftq: Changed 0 to 2 Edge case for <=1 sec videos
441
+ context_frames = history_latents[:, :, -total_context_frames:, :, :] if total_context_frames > 0 else history_latents[:, :, :fallback_frame_count, :, :]
442
+ if total_context_frames > 0:
443
+ split_sizes = [num_4x_frames, num_2x_frames, effective_clean_frames]
444
+ split_sizes = [s for s in split_sizes if s > 0] # Remove zero sizes
445
+ if split_sizes:
446
+ splits = context_frames.split(split_sizes, dim=2)
447
+ split_idx = 0
448
+ clean_latents_4x = splits[split_idx] if num_4x_frames > 0 else history_latents[:, :, :fallback_frame_count, :, :]
449
+ if clean_latents_4x.shape[2] < 2: # 20250507 pftq: edge case for <=1 sec videos
450
+ clean_latents_4x = torch.cat([clean_latents_4x, clean_latents_4x[:, :, -1:, :, :]], dim=2)[:, :, :2, :, :]
451
+ split_idx += 1 if num_4x_frames > 0 else 0
452
+ clean_latents_2x = splits[split_idx] if num_2x_frames > 0 and split_idx < len(splits) else history_latents[:, :, :fallback_frame_count, :, :]
453
+ if clean_latents_2x.shape[2] < 2: # 20250507 pftq: edge case for <=1 sec videos
454
+ clean_latents_2x = torch.cat([clean_latents_2x, clean_latents_2x[:, :, -1:, :, :]], dim=2)[:, :, :2, :, :]
455
+ split_idx += 1 if num_2x_frames > 0 else 0
456
+ clean_latents_1x = splits[split_idx] if effective_clean_frames > 0 and split_idx < len(splits) else history_latents[:, :, :fallback_frame_count, :, :]
457
+ else:
458
+ clean_latents_4x = clean_latents_2x = clean_latents_1x = history_latents[:, :, :fallback_frame_count, :, :]
459
+ else:
460
+ clean_latents_4x = clean_latents_2x = clean_latents_1x = history_latents[:, :, :fallback_frame_count, :, :]
461
+
462
+ clean_latents = torch.cat([start_latent.to(history_latents), clean_latents_1x], dim=2)
463
+
464
+ # 20250507 pftq: Fix for <=1 sec videos.
465
+ max_frames = min(latent_window_size * 4 - 3, history_latents.shape[2] * 4)
466
+
467
+ generated_latents = sample_hunyuan(
468
+ transformer=transformer,
469
+ sampler='unipc',
470
+ width=width,
471
+ height=height,
472
+ frames=max_frames,
473
+ real_guidance_scale=cfg,
474
+ distilled_guidance_scale=gs,
475
+ guidance_rescale=rs,
476
+ num_inference_steps=steps,
477
+ generator=rnd,
478
+ prompt_embeds=llama_vec,
479
+ prompt_embeds_mask=llama_attention_mask,
480
+ prompt_poolers=clip_l_pooler,
481
+ negative_prompt_embeds=llama_vec_n,
482
+ negative_prompt_embeds_mask=llama_attention_mask_n,
483
+ negative_prompt_poolers=clip_l_pooler_n,
484
+ device=gpu,
485
+ dtype=torch.bfloat16,
486
+ image_embeddings=image_encoder_last_hidden_state,
487
+ latent_indices=latent_indices,
488
+ clean_latents=clean_latents,
489
+ clean_latent_indices=clean_latent_indices,
490
+ clean_latents_2x=clean_latents_2x,
491
+ clean_latent_2x_indices=clean_latent_2x_indices,
492
+ clean_latents_4x=clean_latents_4x,
493
+ clean_latent_4x_indices=clean_latent_4x_indices,
494
+ callback=callback,
495
+ )
496
+
497
+ total_generated_latent_frames += int(generated_latents.shape[2])
498
+ history_latents = torch.cat([history_latents, generated_latents.to(history_latents)], dim=2)
499
+
500
+ if not high_vram:
501
+ offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
502
+ load_model_as_complete(vae, target_device=gpu)
503
+
504
+ real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :]
505
+
506
+ if history_pixels is None:
507
+ history_pixels = vae_decode(real_history_latents, vae).cpu()
508
+ else:
509
+ section_latent_frames = latent_window_size * 2
510
+ overlapped_frames = min(latent_window_size * 4 - 3, history_pixels.shape[2])
511
+
512
+ #if section_index == 0:
513
+ #extra_latents = 1 # Add up to 2 extra latent frames for smoother overlap to initial video
514
+ #extra_pixel_frames = extra_latents * 4 # Approx. 4 pixel frames per latent
515
+ #overlapped_frames = min(overlapped_frames + extra_pixel_frames, history_pixels.shape[2], section_latent_frames * 4)
516
+
517
+ current_pixels = vae_decode(real_history_latents[:, :, -section_latent_frames:], vae).cpu()
518
+ history_pixels = soft_append_bcthw(history_pixels, current_pixels, overlapped_frames)
519
+
520
+ if not high_vram:
521
+ unload_complete_models()
522
+
523
+ output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
524
+
525
+ # 20250506 pftq: Use input video FPS for output
526
+ save_bcthw_as_mp4(history_pixels, output_filename, fps=fps, crf=mp4_crf)
527
+ print(f"Latest video saved: {output_filename}")
528
+ # 20250508 pftq: Save prompt to mp4 metadata comments
529
+ set_mp4_comments_imageio_ffmpeg(output_filename, f"Prompt: {prompt} | Negative Prompt: {n_prompt}");
530
+ print(f"Prompt saved to mp4 metadata comments: {output_filename}")
531
+
532
+ # 20250506 pftq: Clean up previous partial files
533
+ if previous_video is not None and os.path.exists(previous_video):
534
+ try:
535
+ os.remove(previous_video)
536
+ print(f"Previous partial video deleted: {previous_video}")
537
+ except Exception as e:
538
+ print(f"Error deleting previous partial video {previous_video}: {e}")
539
+ previous_video = output_filename
540
+
541
+ print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
542
+
543
+ stream.output_queue.push(('file', output_filename))
544
+ except:
545
+ traceback.print_exc()
546
+
547
+ if not high_vram:
548
+ unload_complete_models(
549
+ text_encoder, text_encoder_2, image_encoder, vae, transformer
550
+ )
551
+
552
+ stream.output_queue.push(('end', None))
553
+ return
554
+
555
+ # 20250506 pftq: Modified process to pass clean frame count, etc from video_encode
556
+ def process(input_video, prompt, n_prompt, seed, batch, resolution, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, no_resize, mp4_crf, num_clean_frames, vae_batch):
557
+ global stream, high_vram
558
+ # 20250506 pftq: Updated assertion for video input
559
+ assert input_video is not None, 'No input video!'
560
+
561
+ yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
562
+
563
+ # 20250507 pftq: Even the H100 needs offloading if the video dimensions are 720p or higher
564
+ if high_vram and (no_resize or resolution>640):
565
+ print("Disabling high vram mode due to no resize and/or potentially higher resolution...")
566
+ high_vram = False
567
+ vae.enable_slicing()
568
+ vae.enable_tiling()
569
+ DynamicSwapInstaller.install_model(transformer, device=gpu)
570
+ DynamicSwapInstaller.install_model(text_encoder, device=gpu)
571
+
572
+ # 20250508 pftq: automatically set distilled cfg to 1 if cfg is used
573
+ if cfg > 1:
574
+ gs = 1
575
+
576
+ stream = AsyncStream()
577
+
578
+ # 20250506 pftq: Pass num_clean_frames, vae_batch, etc
579
+ async_run(worker, input_video, prompt, n_prompt, seed, batch, resolution, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, no_resize, mp4_crf, num_clean_frames, vae_batch)
580
+
581
+ output_filename = None
582
+
583
+ while True:
584
+ flag, data = stream.output_queue.next()
585
+
586
+ if flag == 'file':
587
+ output_filename = data
588
+ yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
589
+
590
+ if flag == 'progress':
591
+ preview, desc, html = data
592
+ #yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
593
+ yield output_filename, gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True) # 20250506 pftq: Keep refreshing the video in case it got hidden when the tab was in the background
594
+
595
+ if flag == 'end':
596
+ yield output_filename, gr.update(visible=False), desc+' Video complete.', '', gr.update(interactive=True), gr.update(interactive=False)
597
+ break
598
+
599
+ def end_process():
600
+ stream.input_queue.push('end')
601
+
602
+ quick_prompts = [
603
+ 'The girl dances gracefully, with clear movements, full of charm.',
604
+ 'A character doing some simple body movements.',
605
+ ]
606
+ quick_prompts = [[x] for x in quick_prompts]
607
+
608
+ css = make_progress_bar_css()
609
+ block = gr.Blocks(css=css).queue()
610
+ with block:
611
+ # 20250506 pftq: Updated title to reflect video input functionality
612
+ gr.Markdown('# Framepack F1 with Video Input (Video Extension)')
613
+ with gr.Row():
614
+ with gr.Column():
615
+ # 20250506 pftq: Changed to Video input from Image
616
+ input_video = gr.Video(sources='upload', label="Input Video", height=320)
617
+ prompt = gr.Textbox(label="Prompt", value='')
618
+ #example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Quick List', samples_per_page=1000, components=[prompt])
619
+ #example_quick_prompts.click(lambda x: x[0], inputs=[example_quick_prompts], outputs=prompt, show_progress=False, queue=False)
620
+
621
+ with gr.Row():
622
+ start_button = gr.Button(value="Start Generation")
623
+ end_button = gr.Button(value="End Generation", interactive=False)
624
+
625
+ with gr.Group():
626
+ with gr.Row():
627
+ use_teacache = gr.Checkbox(label='Use TeaCache', value=False, info='Faster speed, but often makes hands and fingers slightly worse.')
628
+ no_resize = gr.Checkbox(label='Force Original Video Resolution (No Resizing)', value=False, info='Might run out of VRAM (720p requires > 24GB VRAM).')
629
+
630
+ seed = gr.Number(label="Seed", value=31337, precision=0)
631
+
632
+ batch = gr.Slider(label="Batch Size (Number of Videos)", minimum=1, maximum=1000, value=1, step=1, info='Generate multiple videos each with a different seed.')
633
+
634
+ resolution = gr.Number(label="Resolution (max width or height)", value=640, precision=0, visible=False)
635
+
636
+ total_second_length = gr.Slider(label="Additional Video Length to Generate (Seconds)", minimum=1, maximum=120, value=5, step=0.1)
637
+
638
+ # 20250506 pftq: Reduced default distilled guidance scale to improve adherence to input video
639
+ gs = gr.Slider(label="Distilled CFG Scale", minimum=1.0, maximum=32.0, value=3.0, step=0.01, info='Prompt adherence at the cost of less details from the input video, but to a lesser extent than Context Frames.')
640
+ cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=32.0, value=1.0, step=0.01, visible=True, info='Use this instead of Distilled for more detail/control + Negative Prompt (make sure Distilled set to 1). Doubles render time.') # Should not change
641
+ rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.01, visible=False) # Should not change
642
+
643
+ n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=True, info='Requires using normal CFG (undistilled) instead of Distilled (set Distilled=1 and CFG > 1).')
644
+ steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1, info='Increase for more quality, especially if using high non-distilled CFG.')
645
+
646
+ # 20250506 pftq: Renamed slider to Number of Context Frames and updated description
647
+ num_clean_frames = gr.Slider(label="Number of Context Frames", minimum=2, maximum=10, value=5, step=1, info="Retain more video details but increase memory use. Reduce to 2 if memory issues.")
648
+
649
+ default_vae = 32
650
+ if high_vram:
651
+ default_vae = 128
652
+ elif free_mem_gb>=20:
653
+ default_vae = 64
654
+
655
+ vae_batch = gr.Slider(label="VAE Batch Size for Input Video", minimum=4, maximum=256, value=default_vae, step=4, info="Reduce if running out of memory. Increase for better quality frames during fast motion.")
656
+
657
+ latent_window_size = gr.Slider(label="Latent Window Size", minimum=9, maximum=33, value=9, step=1, visible=True, info='Generate more frames at a time (larger chunks). Less degradation and better blending but higher VRAM cost.')
658
+
659
+ gpu_memory_preservation = gr.Slider(label="GPU Inference Preserved Memory (GB) (larger means slower)", minimum=6, maximum=128, value=6, step=0.1, info="Set this number to a larger value if you encounter OOM. Larger value causes slower speed.")
660
+
661
+ mp4_crf = gr.Slider(label="MP4 Compression", minimum=0, maximum=100, value=16, step=1, info="Lower means better quality. 0 is uncompressed. Change to 16 if you get black outputs. ")
662
+
663
+ with gr.Column():
664
+ preview_image = gr.Image(label="Next Latents", height=200, visible=False)
665
+ result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=512, loop=True)
666
+ progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
667
+ progress_bar = gr.HTML('', elem_classes='no-generating-animation')
668
+
669
+ gr.HTML("""
670
+ <div style="text-align:center; margin-top:20px;">Share your results and find ideas at the <a href="https://x.com/search?q=framepack&f=live" target="_blank">FramePack Twitter (X) thread</a></div>
671
+ """)
672
+
673
+ # 20250506 pftq: Updated inputs to include num_clean_frames
674
+ ips = [input_video, prompt, n_prompt, seed, batch, resolution, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, no_resize, mp4_crf, num_clean_frames, vae_batch]
675
+ start_button.click(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button])
676
+ end_button.click(fn=end_process)
677
+
678
+ block.launch(
679
+ server_name=args.server,
680
+ server_port=args.port,
681
+ share=args.share,
682
+ inbrowser=args.inbrowser,
683
+ )