app.py

import spaces
import gradio as gr
import safetensors.torch
import torchvision.transforms.v2 as transforms
import cv2
import torch
import numpy as np
from typing import List, Optional, Tuple, Union
from PIL import Image
import io
from io import BytesIO
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers import HunyuanVideoPipeline, FlowMatchEulerDiscreteScheduler
from diffusers.models.transformers.transformer_hunyuan_video import HunyuanVideoPatchEmbed, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video
from diffusers.models.attention import Attention
from diffusers.utils.state_dict_utils import convert_state_dict_to_diffusers, convert_unet_state_dict_to_peft
from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict
from diffusers.models.embeddings import apply_rotary_emb
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.loaders import HunyuanVideoLoraLoaderMixin
from diffusers.models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
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
import requests
import io


# Define video transformations
video_transforms = transforms.Compose(
    [
        transforms.Lambda(lambda x: x / 255.0),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
    ]
)
quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_8bit = HunyuanVideoTransformer3DModel.from_pretrained(
    "hunyuanvideo-community/HunyuanVideo",
    subfolder="transformer",
    quantization_config=quant_config,
    torch_dtype=torch.bfloat16,
)

pipeline = HunyuanVideoPipeline.from_pretrained(
    "hunyuanvideo-community/HunyuanVideo",
    transformer=transformer_8bit,
    torch_dtype=torch.float16,
    device_map="balanced",
)
model_id = "hunyuanvideo-community/HunyuanVideo"
lora_path = hf_hub_download("dashtoon/hunyuan-video-keyframe-control-lora", "i2v.sft")  
# lora_path = "./cache/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)
# pipe.to("cuda")
# Enable memory savings
# pipe.vae.enable_slicing()
pipe.vae.enable_tiling()
pipe.enable_model_cpu_offload()

with torch.no_grad():  # enable image inputs
    initial_input_channels = pipe.transformer.config.in_channels
    new_img_in = HunyuanVideoPatchEmbed(
        patch_size=(pipe.transformer.config.patch_size_t, pipe.transformer.config.patch_size, pipe.transformer.config.patch_size),
        in_chans=pipe.transformer.config.in_channels * 2,
        embed_dim=pipe.transformer.config.num_attention_heads * pipe.transformer.config.attention_head_dim,
    )
    new_img_in = new_img_in.to(pipe.device, dtype=pipe.dtype)
    new_img_in.proj.weight.zero_()
    new_img_in.proj.weight[:, :initial_input_channels].copy_(pipe.transformer.x_embedder.proj.weight)
    if pipe.transformer.x_embedder.proj.bias is not None:
        new_img_in.proj.bias.copy_(pipe.transformer.x_embedder.proj.bias)
    pipe.transformer.x_embedder = new_img_in

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"])
pipe.unload_lora_weights()

# Function to read the content of a markdown file in the same directory
def read_markdown_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read()
        
def resize_image_to_bucket(image: Union[Image.Image, np.ndarray], bucket_reso: Tuple[int, int]) -> np.ndarray:
    """
    Resize the image to the bucket resolution.
    """
    if isinstance(image, Image.Image):
        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
    bucket_width, bucket_height = bucket_reso
    scale_width = bucket_width / image_width
    scale_height = bucket_height / image_height
    scale = max(scale_width, scale_height)
    image_width = int(image_width * scale + 0.5)
    image_height = int(image_height * scale + 0.5)
    if scale > 1:
        image = Image.fromarray(image)
        image = image.resize((image_width, image_height), Image.LANCZOS)
        image = np.array(image)
    else:
        image = cv2.resize(image, (image_width, image_height), interpolation=cv2.INTER_AREA)
    # crop the image to the bucket resolution
    crop_left = (image_width - bucket_width) // 2
    crop_top = (image_height - bucket_height) // 2
    image = image[crop_top:crop_top + bucket_height, crop_left:crop_left + bucket_width]
    return image

# 
# @torch.inference_mode()
@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)
    cond_frame1 = resize_image_to_bucket(cond_frame1, bucket_reso=(width, height))
    cond_frame2 = resize_image_to_bucket(cond_frame2, bucket_reso=(width, height))
    cond_video = np.zeros(shape=(num_frames, height, width, 3))
    cond_video[0], cond_video[-1] = cond_frame1, cond_frame2
    cond_video = torch.from_numpy(cond_video.copy()).permute(0, 3, 1, 2)
    cond_video = torch.stack([video_transforms(x) for x in cond_video], dim=0).unsqueeze(0)
    with torch.no_grad():
        image_or_video = cond_video.to(device="cuda", dtype=pipe.dtype)
        image_or_video = image_or_video.permute(0, 2, 1, 3, 4).contiguous()  # [B, F, C, H, W] -> [B, C, F, H, W]
        cond_latents = pipe.vae.encode(image_or_video).latent_dist.sample()
        cond_latents = cond_latents * pipe.vae.config.scaling_factor
        cond_latents = cond_latents.to(dtype=pipe.dtype)
        assert not torch.any(torch.isnan(cond_latents))
    # Generate video
    video = call_pipe(
        pipe,
        prompt=prompt,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        image_latents=cond_latents,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        generator=torch.Generator(device="cuda").manual_seed(0),
    ).frames[0]
    # Export to video
    # TO-DO: Implement alternate method
    video_path = "output.mp4"
    export_to_video(video, video_path, fps=24)
    torch.cuda.empty_cache()
    return video_path

@torch.inference_mode()
def call_pipe(
    pipe,
    prompt: Union[str, List[str]] = None,
    prompt_2: Union[str, List[str]] = None,
    height: int = 720,
    width: int = 1280,
    num_frames: int = 129,
    num_inference_steps: int = 50,
    sigmas: Optional[List[float]] = None,
    guidance_scale: float = 6.0,
    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,
    pooled_prompt_embeds: Optional[torch.Tensor] = None,
    prompt_attention_mask: Optional[torch.Tensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    attention_kwargs: Optional[dict] = None,
    callback_on_step_end: Optional[Union[callable, PipelineCallback, MultiPipelineCallbacks]] = None,
    callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
    prompt_template: Optional[dict] = DEFAULT_PROMPT_TEMPLATE,
    max_sequence_length: int = 256,
    image_latents: Optional[torch.Tensor] = None,
):
    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,
        prompt_2,
        height,
        width,
        prompt_embeds,
        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
    elif prompt is not None and isinstance(prompt, list):
        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,
        prompt_2=prompt_2,
        prompt_template=prompt_template,
        num_videos_per_prompt=num_videos_per_prompt,
        prompt_embeds=prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        prompt_attention_mask=prompt_attention_mask,
        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(
        pipe.scheduler,
        num_inference_steps,
        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
    latents = pipe.prepare_latents(
        batch_size * num_videos_per_prompt,
        num_channels_latents,
        height,
        width,
        num_latent_frames,
        torch.float32,
        device,
        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_2.to("cpu")  
    torch.cuda.empty_cache()
    with pipe.progress_bar(total=num_inference_steps) as progress_bar:
        for i, t in enumerate(timesteps):
            if pipe.interrupt:
                continue
            pipe._current_timestep = t
            latent_model_input = latents.to(transformer_dtype)
            timestep = t.expand(latents.shape[0]).to(latents.dtype)
            noise_pred = pipe.transformer(
                hidden_states=torch.cat([latent_model_input, image_latents], dim=1),
                timestep=timestep,
                encoder_hidden_states=prompt_embeds,
                encoder_attention_mask=prompt_attention_mask,
                pooled_projections=pooled_prompt_embeds,
                guidance=guidance,
                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:
                    callback_kwargs[k] = locals()[k]
                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):
                progress_bar.update()
    pipe._current_timestep = None
    if not output_type == "latent":
        latents = latents.to(pipe.vae.dtype) / pipe.vae.config.scaling_factor
        video = pipe.vae.decode(latents, return_dict=False)[0]
        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

    with gr.Blocks(css=".gradio-container { max-width: 80vw; margin: 0 auto; }, /* Target all media elements (img, video, audio) within table cells */ tr td img, tr td video, tr td audio { max-height: 240px; object-fit: contain; display: block; width: auto; } /* Target all table cells within table rows */ tr td { overflow: hidden; }") as demo:
        with gr.Group():
            gr.Markdown("""# HunyuanVideo Keyframe Control Lora for Video Generation
                            **Generate videos using the HunyuanVideo model with a prompt and two (or more) frames as conditions. Gradio / HF Spaces implementation demo.**
                            ---
                            For more technical information check out the [original repo by dashtoon.](https://huggingface.co/dashtoon/hunyuan-video-keyframe-control-lora) Special shoutout to @pftq for work on optimization and ideas. Gradio Implementation by [AI Without Borders](https://huggingface.co/aiwithoutborders-xyz); this repo will be moved to the org's namespace once billing is sorted.

                            * Unfortunately, it's still difficult to run on a ZeroGPU space, but we're getting closer. Until then, or until we are granted a GPU allocation, this space was created for you to **DUPLICATE** and begin generating on your own hardware.. 
                            
                            I will fill out a request for GPU allocation for the demo with HF soon.
            
            """)


                            
        with gr.Row():
            with gr.Column(scale=5):
                with gr.Row():
                    prompt_textbox = gr.Textbox(label="Prompt", value="a subject ...", scale=2)
                    resolution = gr.Dropdown(
                        label="Resolution",
                        choices=["720x1280", "544x960", "1280x720", "960x544", "720x720"],
                        value="544x960"
                    )
                frame1 = gr.Image(label="Frame 1", type="pil")
                frame2 = gr.Image(label="Frame 2", type="pil")
                num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=30)
                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)
                generate_button = gr.Button("Generate Video")
            with gr.Column(scale=3):
                outputs = gr.Video(label="Generated Video")
                with gr.Accordion(label="Examples"):
                    markdown_content = read_markdown_file("examples.md")
                    gr.Markdown(markdown_content, sanitize_html=False)
                with gr.Accordion():
                    gr.Markdown("""

                    ## 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.​

                    ## Recommended Settings
                    1. The model works best on human subjects. Single subject images work slightly better.
                    2. It is recommended to use the following image generation resolutions `720x1280`, `544x960`, `1280x720`, `960x544`.
                    3. It is recommended to set frames from 33 upto 97. Can go upto 121 frames as well (but not tested much).
                    4. Prompting helps a lot but works even without. The prompt can be as simple as just the name of the object you want to generate or can be detailed.
                    5. `num_inference_steps` is recommended to be 50, but for fast results you can use 30 as well. Anything less than 30 is not recommended.

                    ## FINAL THOUGHTS: This ZeroGPU space, while successfully loaded, has its memory packed to the rim. If you're lucky you may be able to sneak in a small demo inference here and there, but you will most definitely not be using the recommended settings listed above. Help, of course, is not only welcome but very much appreciated. Learn more about our non-profit initiative, [AI Without Borders](https://huggingface.co/aiwithoutborders-xyz), by following us on Huggingface or on [X](http://x.com/borderlesstools), where we will be announcing a handful of exciting developments. 
                
                """, sanitize_html=False, elem_id="md_footer", container=True)

        generate_button.click(generate_video, inputs=[prompt_textbox, frame1, frame2, resolution, guidance_scale, num_frames, num_inference_steps], outputs=outputs)

    demo.launch(show_error=True)

if __name__ == "__main__":
    main()