Merge

app_start_end.py

@@ -313,7 +313,277 @@ def set_mp4_comments_imageio_ffmpeg(input_file, comments):
         return False
 
 @torch.no_grad()
-def worker_start_end(input_image, end_image, prompts, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, fps_number):
+def worker(input_image, image_position, end_image, prompts, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, fps_number):
+    def encode_prompt(prompt, n_prompt):
+        llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
+
+        if cfg == 1:
+            llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
+        else:
+            llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
+
+        llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
+        llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
+
+        llama_vec = llama_vec.to(transformer.dtype)
+        llama_vec_n = llama_vec_n.to(transformer.dtype)
+        clip_l_pooler = clip_l_pooler.to(transformer.dtype)
+        clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
+        return [llama_vec, clip_l_pooler, llama_vec_n, clip_l_pooler_n, llama_attention_mask, llama_attention_mask_n]
+
+    total_latent_sections = (total_second_length * fps_number) / (latent_window_size * 4)
+    total_latent_sections = int(max(round(total_latent_sections), 1))
+
+    first_section_index = max(min(math.floor(image_position * (total_latent_sections - 1) / 100), (total_latent_sections - 1)), 0)
+    section_index = first_section_index
+    forward = (image_position == 0)
+
+    job_id = generate_timestamp()
+
+    stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
+
+    try:
+        # Clean GPU
+        if not high_vram:
+            unload_complete_models(
+                text_encoder, text_encoder_2, image_encoder, vae, transformer
+            )
+
+        # Text encoding
+
+        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
+
+        if not high_vram:
+            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.
+            load_model_as_complete(text_encoder_2, target_device=gpu)
+
+        prompt_parameters = []
+
+        for prompt_part in prompts[:total_latent_sections]:
+            prompt_parameters.append(encode_prompt(prompt_part, n_prompt))
+
+        # Clean GPU
+        if not high_vram:
+            unload_complete_models(
+                text_encoder, text_encoder_2
+            )
+
+        # Processing input image
+
+        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
+
+        H, W, C = input_image.shape
+        height, width = find_nearest_bucket(H, W, resolution=resolution)
+        
+        def get_start_latent(input_image, height, width, vae, gpu, image_encoder, high_vram):
+            input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
+            
+            #Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
+            
+            input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
+            input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
+            
+            # VAE encoding
+            
+            stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
+            
+            if not high_vram:
+                load_model_as_complete(vae, target_device=gpu)
+            
+            start_latent = vae_encode(input_image_pt, vae)
+            
+            # CLIP Vision
+            
+            stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
+            
+            if not high_vram:
+                unload_complete_models(vae)
+                load_model_as_complete(image_encoder, target_device=gpu)
+
+            image_encoder_last_hidden_state = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder).last_hidden_state
+            
+            if not high_vram:
+                unload_complete_models(image_encoder)
+                
+            return [start_latent, image_encoder_last_hidden_state]
+            
+        [start_latent, image_encoder_last_hidden_state] = get_start_latent(input_image, height, width, vae, gpu, image_encoder, high_vram)
+
+        # Dtype
+
+        image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
+
+        # Sampling
+
+        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
+
+        rnd = torch.Generator("cpu").manual_seed(seed)
+
+        history_latents = torch.zeros(size=(1, 16, 16 + 2 + 1, height // 8, width // 8), dtype=torch.float32, device=cpu)
+        start_latent = start_latent.to(history_latents)
+        history_pixels = None
+
+        history_latents = torch.cat([history_latents, start_latent] if forward else [start_latent, history_latents], dim=2)
+        total_generated_latent_frames = 1
+
+        if enable_preview:
+            def callback(d):
+                preview = d['denoised']
+                preview = vae_decode_fake(preview)
+			    
+                preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
+                preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
+			    
+                if stream.input_queue.top() == 'end':
+                    stream.output_queue.push(('end', None))
+                    raise KeyboardInterrupt('User ends the task.')
+			    
+                current_step = d['i'] + 1
+                percentage = int(100.0 * current_step / steps)
+                hint = f'Sampling {current_step}/{steps}'
+                desc = f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, Video length: {max(0, (total_generated_latent_frames * 4 - 3) / fps_number) :.2f} seconds (FPS-30), Resolution: {height}px * {width}px. The video is being extended now ...'
+                stream.output_queue.push(('progress', (preview, desc, make_progress_bar_html(percentage, hint))))
+                return
+        else:
+            def callback(d):
+                return
+
+        indices = torch.arange(0, 1 + 16 + 2 + 1 + latent_window_size).unsqueeze(0)
+        if forward:
+            clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
+            clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
+        else:
+            latent_indices, clean_latent_1x_indices, clean_latent_2x_indices, clean_latent_4x_indices, clean_latent_indices_start = indices.split([latent_window_size, 1, 2, 16, 1], dim=1)
+            clean_latent_indices = torch.cat([clean_latent_1x_indices, clean_latent_indices_start], dim=1)
+
+        def post_process(forward, generated_latents, total_generated_latent_frames, history_latents, high_vram, transformer, gpu, vae, history_pixels, latent_window_size, enable_preview, section_index, total_latent_sections, outputs_folder, mp4_crf, stream):
+            total_generated_latent_frames += int(generated_latents.shape[2])
+            history_latents = torch.cat([history_latents, generated_latents.to(history_latents)] if forward else [generated_latents.to(history_latents), history_latents], dim=2)
+
+            if not high_vram:
+                offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
+                load_model_as_complete(vae, target_device=gpu)
+
+            if history_pixels is None:
+                real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :] if forward else history_latents[:, :, :total_generated_latent_frames, :, :]
+                history_pixels = vae_decode(real_history_latents, vae).cpu()
+            else:
+                section_latent_frames = latent_window_size * 2
+                overlapped_frames = latent_window_size * 4 - 3
+
+                if forward:
+                    real_history_latents = history_latents[:, :, -min(section_latent_frames, total_generated_latent_frames):, :, :]
+                    history_pixels = soft_append_bcthw(history_pixels, vae_decode(real_history_latents, vae).cpu(), overlapped_frames)
+                else:
+                    real_history_latents = history_latents[:, :, :min(section_latent_frames, total_generated_latent_frames), :, :]
+                    history_pixels = soft_append_bcthw(vae_decode(real_history_latents, vae).cpu(), history_pixels, overlapped_frames)
+
+            if not high_vram:
+                unload_complete_models(text_encoder, text_encoder_2, image_encoder, vae, transformer)
+
+            if enable_preview or section_index == (0 if first_section_index == (total_latent_sections - 1) else (total_latent_sections - 1)):
+                output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
+
+                save_bcthw_as_mp4(history_pixels, output_filename, fps=fps_number, crf=mp4_crf)
+
+                print(f'Decoded. Current latent shape pixel shape {history_pixels.shape}')
+
+                stream.output_queue.push(('file', output_filename))
+            return [total_generated_latent_frames, history_latents, history_pixels]
+        
+        while section_index < total_latent_sections:
+            if stream.input_queue.top() == 'end':
+                stream.output_queue.push(('end', None))
+                return
+
+            print(f'section_index = {section_index}, total_latent_sections = {total_latent_sections}')
+
+            prompt_index = min(section_index, len(prompt_parameters) - 1)
+
+            [llama_vec, clip_l_pooler, llama_vec_n, clip_l_pooler_n, llama_attention_mask, llama_attention_mask_n] = prompt_parameters[prompt_index]
+            
+            if prompt_index < len(prompt_parameters) - 1 or (prompt_index == total_latent_sections - 1):
+                prompt_parameters[prompt_index] = None
+
+            if not high_vram:
+                unload_complete_models()
+                move_model_to_device_with_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
+
+            if use_teacache:
+                transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
+            else:
+                transformer.initialize_teacache(enable_teacache=False)
+
+            if forward:
+                clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[:, :, -(16 + 2 + 1):, :, :].split([16, 2, 1], dim=2)
+                clean_latents = torch.cat([start_latent, clean_latents_1x], dim=2)
+            else:
+                clean_latents_1x, clean_latents_2x, clean_latents_4x = history_latents[:, :, :(1 + 2 + 16), :, :].split([1, 2, 16], dim=2)
+                clean_latents = torch.cat([clean_latents_1x, start_latent], dim=2)
+
+            generated_latents = sample_hunyuan(
+                transformer=transformer,
+                sampler='unipc',
+                width=width,
+                height=height,
+                frames=latent_window_size * 4 - 3,
+                real_guidance_scale=cfg,
+                distilled_guidance_scale=gs,
+                guidance_rescale=rs,
+                # shift=3.0,
+                num_inference_steps=steps,
+                generator=rnd,
+                prompt_embeds=llama_vec,
+                prompt_embeds_mask=llama_attention_mask,
+                prompt_poolers=clip_l_pooler,
+                negative_prompt_embeds=llama_vec_n,
+                negative_prompt_embeds_mask=llama_attention_mask_n,
+                negative_prompt_poolers=clip_l_pooler_n,
+                device=gpu,
+                dtype=torch.bfloat16,
+                image_embeddings=image_encoder_last_hidden_state,
+                latent_indices=latent_indices,
+                clean_latents=clean_latents,
+                clean_latent_indices=clean_latent_indices,
+                clean_latents_2x=clean_latents_2x,
+                clean_latent_2x_indices=clean_latent_2x_indices,
+                clean_latents_4x=clean_latents_4x,
+                clean_latent_4x_indices=clean_latent_4x_indices,
+                callback=callback,
+            )
+
+            [total_generated_latent_frames, history_latents, history_pixels] = post_process(forward, generated_latents, total_generated_latent_frames, history_latents, high_vram, transformer, gpu, vae, history_pixels, latent_window_size, enable_preview, section_index, total_latent_sections, outputs_folder, mp4_crf, stream)
+            
+            if not forward:
+                if section_index > 0:
+                    section_index -= 1
+                else:
+                    clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
+                    clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
+
+                    real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :]
+                    zero_latents = history_latents[:, :, total_generated_latent_frames:, :, :]
+                    history_latents = torch.cat([zero_latents, real_history_latents], dim=2)
+                    real_history_latents = zero_latents = None
+                    
+                    forward = True
+                    section_index = first_section_index
+            
+            if forward:
+                section_index += 1
+    except:
+        traceback.print_exc()
+
+        if not high_vram:
+            unload_complete_models(
+                text_encoder, text_encoder_2, image_encoder, vae, transformer
+            )
+
+    stream.output_queue.push(('end', None))
+    return
+
+@torch.no_grad()
+def worker_start_end(input_image, image_position, end_image, prompts, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, fps_number):
     def encode_prompt(prompt, n_prompt):
         llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
 
@@ -855,18 +1125,18 @@ def worker_video(input_video, prompts, n_prompt, seed, batch, resolution, total_
     stream.output_queue.push(('end', None))
     return
 
-def get_duration_start_end(input_image, image_position, end_image, prompts, generation_mode, n_prompt, seed, resolution, total_second_length, allocation_time, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, enable_preview, use_teacache, mp4_crf, fps_number):
+def get_duration(input_image, image_position, end_image, prompts, generation_mode, n_prompt, seed, resolution, total_second_length, allocation_time, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, enable_preview, use_teacache, mp4_crf, fps_number):
     return allocation_time
 
 # Remove this decorator if you run on local
-@spaces.GPU(duration=get_duration_start_end)
-def process_start_end_on_gpu(input_image, image_position, end_image, prompts, generation_mode, n_prompt, seed, resolution, total_second_length, allocation_time, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, enable_preview, use_teacache, mp4_crf, fps_number
+@spaces.GPU(duration=get_duration)
+def process_on_gpu(input_image, image_position, end_image, prompts, generation_mode, n_prompt, seed, resolution, total_second_length, allocation_time, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, enable_preview, use_teacache, mp4_crf, fps_number
            ):
     start = time.time()
     global stream
     stream = AsyncStream()
 
-    async_run(worker_start_end, input_image, end_image, prompts, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, fps_number)
+    async_run(worker_start_end if generation_mode == "start_end" else worker, input_image, image_position, end_image, prompts, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf, fps_number)
 
     output_filename = None
 
@@ -895,7 +1165,7 @@ def process_start_end_on_gpu(input_image, image_position, end_image, prompts, ge
             "You can upscale the result with RIFE. To make all your generated scenes consistent, you can then apply a face swap on the main character. If you do not see the generated video above, the process may have failed. See the logs for more information. If you see an error like ''NVML_SUCCESS == r INTERNAL ASSERT FAILED'', you probably haven't enough VRAM. Test an example or other options to compare. You can share your inputs to the original space or set your space in public for a peer review.", gr.update(interactive=True), gr.update(interactive=False), gr.update(visible = False)
             break
 
-def process_start_end(input_image,
+def process(input_image,
             image_position=0,
             end_image=None,
             prompt="",
@@ -947,7 +1217,7 @@ def process_start_end(input_image,
 
     yield gr.update(label="Previewed Frames"), None, '', '', gr.update(interactive=False), gr.update(interactive=True), gr.skip()
 
-    yield from process_start_end_on_gpu(input_image,
+    yield from process_on_gpu(input_image,
             image_position,
             end_image,
             prompts,
@@ -1120,7 +1390,7 @@ with block:
     local_storage = gr.BrowserState(default_local_storage)
     with gr.Row():
         with gr.Column():
-            generation_mode = gr.Radio([["Text-to-Video", "text"], ["Image-to-Video", "image"], ["Video Extension", "video"]], elem_id="generation-mode", label="Generation mode", value = "image")
+            generation_mode = gr.Radio([["Text-to-Video", "text"], ["Image-to-Video", "image"], ["Start frame & End frame", "start_end"], ["Video Extension", "video"]], elem_id="generation-mode", label="Generation mode", value = "image")
             text_to_video_hint = gr.HTML("Text-to-Video badly works with a flash effect at the start. I discourage to use the Text-to-Video feature. You should rather generate an image with Flux and use Image-to-Video. You will save time.")
             input_image = gr.Image(sources='upload', type="numpy", label="Image", height=320)
             end_image = gr.Image(sources='upload', type="numpy", label="End Frame (Optional)", height=320)
@@ -1244,7 +1514,7 @@ with block:
                     ]
                 ],
             run_on_click = True,
-            fn = process_start_end,
+            fn = process,
     	    inputs = ips,
             outputs = [result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning],
             cache_examples = torch.cuda.device_count() > 0,
@@ -1376,7 +1646,43 @@ with block:
                     ],
                 ],
             run_on_click = True,
-            fn = process_start_end,
+            fn = process,
+	        inputs = ips,
+            outputs = [result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning],
+            cache_examples = torch.cuda.device_count() > 0,
+        )
+
+    with gr.Row(elem_id="start_end_examples", visible=False):
+        gr.Examples(
+        label = "Examples from start and end frames",
+            examples = [
+                    [
+                        "./img_examples/Example2.webp", # input_image
+                        0, # image_position
+                        None, # end_image
+                        "A man on the left and a woman on the right face each other ready to start a conversation, large space between the persons, full view, full-length view, 3D, pixar, 3D render, CGI. The man talks and the woman listens; A man on the left and a woman on the right face each other ready to start a conversation, large space between the persons, full view, full-length view, 3D, pixar, 3D render, CGI. The woman talks, the man stops talking and the man listens; A man on the left and a woman on the right face each other ready to start a conversation, large space between the persons, full view, full-length view, 3D, pixar, 3D render, CGI. The woman talks and the man listens",
+                        "start_end", # generation_mode
+                        "Missing arm, long hand, unrealistic position, impossible contortion, visible bone, muscle contraction, poorly framed, blurred, blurry, over-smooth", # n_prompt
+                        True, # randomize_seed
+                        42, # seed
+                        True, # auto_allocation
+                        180, # allocation_time
+                        672, # resolution
+                        1, # total_second_length
+                        9, # latent_window_size
+                        30, # steps
+                        1.0, # cfg
+                        10.0, # gs
+                        0.0, # rs
+                        6, # gpu_memory_preservation
+                        False, # enable_preview
+                        False, # use_teacache
+                        16, # mp4_crf
+                        30 # fps_number
+                    ],
+                ],
+            run_on_click = True,
+            fn = process,
 	        inputs = ips,
             outputs = [result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning],
             cache_examples = torch.cuda.device_count() > 0,
@@ -1471,7 +1777,7 @@ with block:
                 ]
             ],
         run_on_click = True,
-        fn = process_start_end,
+        fn = process,
 	    inputs = ips,
         outputs = [result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning],
         cache_examples = False,
@@ -1602,7 +1908,42 @@ with block:
                 ]
             ],
         run_on_click = True,
-        fn = process_start_end,
+        fn = process,
+	    inputs = ips,
+        outputs = [result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning],
+        cache_examples = False,
+    )
+
+    gr.Examples(
+        label = "🖼️ Examples from start and end frames",
+        examples = [
+                [
+                    "./img_examples/Example1.png", # input_image
+                    0, # image_position
+                    None, # end_image
+                    "A dolphin emerges from the water, photorealistic, realistic, intricate details, 8k, insanely detailed",
+                    "start_end", # generation_mode
+                    "Missing arm, long hand, unrealistic position, impossible contortion, visible bone, muscle contraction, poorly framed, blurred, blurry, over-smooth", # n_prompt
+                    True, # randomize_seed
+                    42, # seed
+                    True, # auto_allocation
+                    180, # allocation_time
+                    672, # resolution
+                    1, # total_second_length
+                    9, # latent_window_size
+                    30, # steps
+                    1.0, # cfg
+                    10.0, # gs
+                    0.0, # rs
+                    6, # gpu_memory_preservation
+                    False, # enable_preview
+                    True, # use_teacache
+                    16, # mp4_crf
+                    30 # fps_number
+                ],
+            ],
+        run_on_click = True,
+        fn = process,
 	    inputs = ips,
         outputs = [result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning],
         cache_examples = False,
@@ -1714,7 +2055,7 @@ with block:
     timeless_prompt.change(fn=handle_timeless_prompt_change, inputs=[timeless_prompt], outputs=[final_prompt])
     start_button.click(fn = check_parameters, inputs = [
         generation_mode, input_image, input_video
-    ], outputs = [end_button, warning], queue = False, show_progress = False).success(fn=process_start_end, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning], scroll_to_output = True)
+    ], outputs = [end_button, warning], queue = False, show_progress = False).success(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button, warning], scroll_to_output = True)
     start_button_video.click(fn = check_parameters, inputs = [
         generation_mode, input_image, input_video
     ], outputs = [end_button, warning], queue = False, show_progress = False).success(fn=process_video, inputs=ips_video, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button_video, end_button, warning], scroll_to_output = True)