Update app.py

app.py

@@ -27,7 +27,7 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.hunyuan_video.pipeline_output import HunyuanVideoPipelineOutput
 from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import retrieve_timesteps, DEFAULT_PROMPT_TEMPLATE
 from diffusers.utils import load_image
-from huggingface_hub import hf_hub_download 
+from huggingface_hub import hf_hub_download
 import requests
 import io
 
@@ -38,8 +38,10 @@ video_transforms = transforms.Compose(
         transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
     ]
 )
+
 model_id = "hunyuanvideo-community/HunyuanVideo"
 lora_path = hf_hub_download("dashtoon/hunyuan-video-keyframe-control-lora", "i2v.sft")  # Replace with the actual LORA path
+
 transformer = HunyuanVideoTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16)
 global pipe
 pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.bfloat16)
@@ -64,6 +66,7 @@ with torch.no_grad():  # enable image inputs
 
 lora_state_dict = safetensors.torch.load_file(lora_path, device="cpu")
 transformer_lora_state_dict = {f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.") and "lora" in k}
+
 pipe.load_lora_into_transformer(transformer_lora_state_dict, transformer=pipe.transformer, adapter_name="i2v", _pipeline=pipe)
 pipe.set_adapters(["i2v"], adapter_weights=[1.0])
 pipe.fuse_lora(components=["transformer"], lora_scale=1.0, adapter_names=["i2v"])
@@ -77,7 +80,6 @@ def resize_image_to_bucket(image: Union[Image.Image, np.ndarray], bucket_reso: T
         image = np.array(image)
     elif not isinstance(image, np.ndarray):
         raise ValueError("Image must be a PIL Image or NumPy array")
-
     image_height, image_width = image.shape[:2]
     if bucket_reso == (image_width, image_height):
         return image
@@ -99,17 +101,14 @@ def resize_image_to_bucket(image: Union[Image.Image, np.ndarray], bucket_reso: T
     image = image[crop_top:crop_top + bucket_height, crop_left:crop_left + bucket_width]
     return image
 
-
 @spaces.GPU(duration=120)
 def generate_video(prompt: str, frame1: Image.Image, frame2: Image.Image, resolution: str, guidance_scale: float, num_frames: int, num_inference_steps: int) -> bytes:
     # Debugging print statements
     print(f"Frame 1 Type: {type(frame1)}")
     print(f"Frame 2 Type: {type(frame2)}")
     print(f"Resolution: {resolution}")
-    
     # Parse resolution
     width, height = map(int, resolution.split('x'))
-    
     # Load and preprocess frames
     cond_frame1 = np.array(frame1)
     cond_frame2 = np.array(frame2)
@@ -139,12 +138,12 @@ def generate_video(prompt: str, frame1: Image.Image, frame2: Image.Image, resolu
         generator=torch.Generator(device="cuda").manual_seed(0),
     ).frames[0]
     # Export to video
+    # TO-DO: Implement alternate method
     video_path = "output.mp4"
-    # video_bytes = io.BytesIO()
     export_to_video(video, video_path, fps=24)
     torch.cuda.empty_cache()
     return video_path
-    
+
 @torch.inference_mode()
 def call_pipe(
     pipe,
@@ -159,7 +158,7 @@ def call_pipe(
     num_videos_per_prompt: Optional[int] = 1,
     generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
     latents: Optional[torch.Tensor] = None,
-    prompt_embeds: Optional[torch.Tensor] = None,
+    prompt_embeds<|im_start|>output> Optional[torch.Tensor] = None,
     pooled_prompt_embeds: Optional[torch.Tensor] = None,
     prompt_attention_mask: Optional[torch.Tensor] = None,
     output_type: Optional[str] = "pil",
@@ -173,7 +172,6 @@ def call_pipe(
 ):
     if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
         callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-    
     # 1. Check inputs. Raise error if not correct
     pipe.check_inputs(
         prompt,
@@ -184,13 +182,11 @@ def call_pipe(
         callback_on_step_end_tensor_inputs,
         prompt_template,
     )
-    
     pipe._guidance_scale = guidance_scale
     pipe._attention_kwargs = attention_kwargs
     pipe._current_timestep = None
     pipe._interrupt = False
     device = pipe._execution_device
-    
     # 2. Define call parameters
     if prompt is not None and isinstance(prompt, str):
         batch_size = 1
@@ -198,7 +194,6 @@ def call_pipe(
         batch_size = len(prompt)
     else:
         batch_size = prompt_embeds.shape[0]
-    
     # 3. Encode input prompt
     prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = pipe.encode_prompt(
         prompt=prompt,
@@ -211,13 +206,11 @@ def call_pipe(
         device=device,
         max_sequence_length=max_sequence_length,
     )
-    
     transformer_dtype = pipe.transformer.dtype
     prompt_embeds = prompt_embeds.to(transformer_dtype)
     prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)
     if pooled_prompt_embeds is not None:
         pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype)
-    
     # 4. Prepare timesteps
     sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas
     timesteps, num_inference_steps = retrieve_timesteps(
@@ -226,7 +219,6 @@ def call_pipe(
         device,
         sigmas=sigmas,
     )
-    
     # 5. Prepare latent variables
     num_channels_latents = pipe.transformer.config.in_channels
     num_latent_frames = (num_frames - 1) // pipe.vae_scale_factor_temporal + 1
@@ -241,17 +233,14 @@ def call_pipe(
         generator,
         latents,
     )
-    
     # 6. Prepare guidance condition
     guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
-    
     # 7. Denoising loop
     num_warmup_steps = len(timesteps) - num_inference_steps * pipe.scheduler.order
     pipe._num_timesteps = len(timesteps)
-    pipe.text_encoder.to("cpu") 
+    pipe.text_encoder.to("cpu")
     pipe.text_encoder_2.to("cpu")  
-    torch.cuda.empty_cache() 
-    
+    torch.cuda.empty_cache()
     with pipe.progress_bar(total=num_inference_steps) as progress_bar:
         for i, t in enumerate(timesteps):
             if pipe.interrupt:
@@ -269,10 +258,8 @@ def call_pipe(
                 attention_kwargs=attention_kwargs,
                 return_dict=False,
             )[0]
-            
             # compute the previous noisy sample x_t -> x_t-1
             latents = pipe.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-            
             if callback_on_step_end is not None:
                 callback_kwargs = {}
                 for k in callback_on_step_end_tensor_inputs:
@@ -280,11 +267,9 @@ def call_pipe(
                 callback_outputs = callback_on_step_end(pipe, i, t, callback_kwargs)
                 latents = callback_outputs.pop("latents", latents)
                 prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-            
             # call the callback, if provided
-            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipe.scheduler.order == 0):
+            if i < len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipe.scheduler.order == 0):
                 progress_bar.update()
-    
     pipe._current_timestep = None
     if not output_type == "latent":
         latents = latents.to(pipe.vae.dtype) / pipe.vae.config.scaling_factor
@@ -292,48 +277,69 @@ def call_pipe(
         video = pipe.video_processor.postprocess_video(video, output_type=output_type)
     else:
         video = latents
-    
     # Offload all models
     pipe.maybe_free_model_hooks()
-    
     if not return_dict:
         return (video,)
     return HunyuanVideoPipelineOutput(frames=video)
 
-
 def main():
     # Define the interface inputs
-    inputs = [
-        gr.Textbox(label="Prompt", value="a woman"),
-        gr.Image(label="Frame 1", type="pil"),
-        gr.Image(label="Frame 2", type="pil"),
-        gr.Dropdown(
-            label="Resolution",
-            choices=["720x1280", "544x960", "1280x720", "960x544", "720x720"],
-            value="544x960"
-        ),
-        gr.Slider(minimum=0.1, maximum=20, step=0.1, label="Guidance Scale", value=6.0),
-        gr.Slider(minimum=1, maximum=129, step=1, label="Number of Frames", value=49),
-        gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=30)
-    ]
-    
+    prompt_textbox = gr.Textbox(label="Prompt", value="a subject ... ")
+    frame1 = gr.Image(label="Frame 1", type="pil")
+    frame2 = gr.Image(label="Frame 2", type="pil")
+    resolution = gr.Dropdown(
+        label="Resolution",
+        choices=["720x1280", "544x960", "1280x720", "960x544", "720x720"],
+        value="544x960"
+    )
+    guidance_scale = gr.Slider(minimum=0.1, maximum=20, step=0.1, label="Guidance Scale", value=6.0)
+    num_frames = gr.Slider(minimum=1, maximum=129, step=1, label="Number of Frames", value=49)
+    num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=30)
+
     # Define the interface outputs
-    outputs = [
-        gr.Video(label="Generated Video"),
-    ]
+    outputs = gr.Video(label="Generated Video")
 
-    
-    # Create the Gradio interface
-    iface = gr.Interface(
-        fn=generate_video,
-        inputs=inputs,
-        outputs=outputs,
-        title="HunyuanVideo Keyframe IMG+IMG2VID Control Lora",
-        description="Generate videos using the HunyuanVideo model with a prompt and two frames as conditions. Gradio / HF Spaces implementation demo.",
-    )
-    
-    # Launch the Gradio app
-    iface.launch(show_error=True)
+    with gr.Blocks() as demo:
+        gr.Markdown("# HunyuanVideo Keyframes Gen Control Lora")
+        gr.Markdown("Generate videos using the HunyuanVideo model with a prompt and two (or more) frames as conditions. Gradio / HF Spaces implementation demo.")
+        gr.Markdown("Unfortunately still difficult to run on ZeroGPU spaces, but getting closer. **DUPLICATE THE SPACE** and select hardware with more VRAM. I will fill out request for GPU allocation for the demo with HF shortly.")
+        gr.Markdown("For more technical information check out the [original repo by dashtoon.](dashtoon/hunyuan-video-keyframe-control-lora) Special shoutout to @pftq for work on optimization and ideas.")
+
+        with gr.Row():
+            with gr.Column():
+                prompt_textbox
+                resolution
+            with gr.Column():
+                guided_scale
+                num_frames
+                
+        with gr.Row():
+            frame1
+            frame2
+            
+        with gr.Row():
+            num_inference_steps
+            outputs
+            
+        generate_button = gr.Button("Generate Video")
+        generate_button.click(generate_video, inputs=[prompt_textbox, frame1, frame2, resolution, guidance_scale, num_frames, num_inference_steps], outputs=outputs)
+
+        gr.Markdown("""
+            @dashtoon: HunyuanVideo Keyframe Control Lora is an adapter for HunyuanVideo T2V model for keyframe-based video generation. ​Our architecture builds upon existing models, introducing key enhancements to optimize keyframe-based video generation:​
+            *  We modify the input patch embedding projection layer to effectively incorporate keyframe information. By adjusting the convolutional input parameters, we enable the model to process image inputs within the Diffusion Transformer (DiT) framework.​
+            *  We apply Low-Rank Adaptation (LoRA) across all linear layers and the convolutional input layer. This approach facilitates efficient fine-tuning by introducing low-rank matrices that approximate the weight updates, thereby preserving the base model's foundational capabilities while reducing the number of trainable parameters.
+            * The model is conditioned on user-defined keyframes, allowing precise control over the generated video's start and end frames. This conditioning ensures that the generated content aligns seamlessly with the specified keyframes, enhancing the coherence and narrative flow of the video.​
+            
+            | Image 1 | Image 2 | Generated Video |
+            |---------|---------|-----------------|
+            | ![Image 1](https://content.dashtoon.ai/stability-images/41aeca63-064a-4003-8c8b-bfe2cc80d275.png) | ![Image 2](https://content.dashtoon.ai/stability-images/28956177-3455-4b56-bb6c-73eacef323ca.png) | <video controls autoplay src="https://content.dashtoon.ai/stability-images/14b7dd1a-1f46-4c4c-b4ec-9d0f948712af.mp4"></video> |
+            | ![Image 1](https://content.dashtoon.ai/stability-images/ddabbf2f-4218-497b-8239-b7b882d93000.png) | ![Image 2](https://content.dashtoon.ai/stability-images/b603acba-40a4-44ba-aa26-ed79403df580.png) | <video controls autoplay src="https://content.dashtoon.ai/stability-images/b00ba193-b3b7-41a1-9bc1-9fdaceba6efa.mp4"></video> |
+            | ![Image 1](https://content.dashtoon.ai/stability-images/5298cf0c-0955-4568-935a-2fb66045f21d.png) | ![Image 2](https://content.dashtoon.ai/stability-images/722a4ea7-7092-4323-8e83-3f627e8fd7f8.png) | <video controls autoplay src="https://content.dashtoon.ai/stability-images/0cb84780-4fdf-4ecc-ab48-12e7e1055a39.mp4"></video> |
+            | ![Image 1](https://content.dashtoon.ai/stability-images/69d9a49f-95c0-4e85-bd49-14a039373c8b.png) | ![Image 2](https://content.dashtoon.ai/stability-images/0cef7fa9-e15a-48ec-9bd3-c61921181802.png) | <video controls autoplay src="https://content.dashtoon.ai/stability-images/ce12156f-0ac2-4d16-b489-37e85c61b5b2.mp4"></video> |
+        """)
+
+    demo.launch(show_error=True)
 
 if __name__ == "__main__":
     main()