LPX55/HunYuan-Keyframe2VID-Control-Lora
9
1import spaces2import gradio as gr3import safetensors.torch4import torchvision.transforms.v2 as transforms5import cv26import torch7import numpy as np8from typing import List, Optional, Tuple, Union9from PIL import Image10import io11from io import BytesIO12from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig13from diffusers import HunyuanVideoPipeline, FlowMatchEulerDiscreteScheduler14from diffusers.models.transformers.transformer_hunyuan_video import HunyuanVideoPatchEmbed, HunyuanVideoTransformer3DModel15from diffusers.utils import export_to_video16from diffusers.models.attention import Attention17from diffusers.utils.state_dict_utils import convert_state_dict_to_diffusers, convert_unet_state_dict_to_peft18from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict19from diffusers.models.embeddings import apply_rotary_emb20from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback21from diffusers.loaders import HunyuanVideoLoraLoaderMixin22from diffusers.models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel23from diffusers.schedulers import FlowMatchEulerDiscreteScheduler24from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring25from diffusers.utils.torch_utils import randn_tensor26from diffusers.video_processor import VideoProcessor27from diffusers.pipelines.pipeline_utils import DiffusionPipeline28from diffusers.pipelines.hunyuan_video.pipeline_output import HunyuanVideoPipelineOutput29from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import retrieve_timesteps, DEFAULT_PROMPT_TEMPLATE30from diffusers.utils import load_image31from huggingface_hub import hf_hub_download32import requests33import io34 35 36# Define video transformations37video_transforms = transforms.Compose(38 [39 transforms.Lambda(lambda x: x / 255.0),40 transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),41 ]42)43quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)44transformer_8bit = HunyuanVideoTransformer3DModel.from_pretrained(45 "hunyuanvideo-community/HunyuanVideo",46 subfolder="transformer",47 quantization_config=quant_config,48 torch_dtype=torch.bfloat16,49)50 51pipeline = HunyuanVideoPipeline.from_pretrained(52 "hunyuanvideo-community/HunyuanVideo",53 # transformer=transformer_8bit,54 torch_dtype=torch.float16,55 # device_map="balanced",56)57model_id = "hunyuanvideo-community/HunyuanVideo"58# lora_path = hf_hub_download("dashtoon/hunyuan-video-keyframe-control-lora", "i2v.sft") 59lora_path = "i2v.sft"60 61# Replace with the actual LORA path62 63transformer = HunyuanVideoTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16)64global pipe65pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.bfloat16)66# pipe.to("cuda")67# Enable memory savings68# pipe.vae.enable_slicing()69pipe.vae.enable_tiling()70pipe.enable_model_cpu_offload()71 72with torch.no_grad(): # enable image inputs73 initial_input_channels = pipe.transformer.config.in_channels74 new_img_in = HunyuanVideoPatchEmbed(75 patch_size=(pipe.transformer.config.patch_size_t, pipe.transformer.config.patch_size, pipe.transformer.config.patch_size),76 in_chans=pipe.transformer.config.in_channels * 2,77 embed_dim=pipe.transformer.config.num_attention_heads * pipe.transformer.config.attention_head_dim,78 )79 new_img_in = new_img_in.to(pipe.device, dtype=pipe.dtype)80 new_img_in.proj.weight.zero_()81 new_img_in.proj.weight[:, :initial_input_channels].copy_(pipe.transformer.x_embedder.proj.weight)82 if pipe.transformer.x_embedder.proj.bias is not None:83 new_img_in.proj.bias.copy_(pipe.transformer.x_embedder.proj.bias)84 pipe.transformer.x_embedder = new_img_in85 86lora_state_dict = safetensors.torch.load_file(lora_path)87transformer_lora_state_dict = {f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.") and "lora" in k}88 89pipe.load_lora_into_transformer(transformer_lora_state_dict, transformer=pipe.transformer, adapter_name="i2v", _pipeline=pipe)90pipe.set_adapters(["i2v"], adapter_weights=[1.0])91pipe.fuse_lora(components=["transformer"], lora_scale=1.0, adapter_names=["i2v"])92pipe.unload_lora_weights()93 94# Function to read the content of a markdown file in the same directory95def read_markdown_file(file_path):96 with open(file_path, 'r', encoding='utf-8') as file:97 return file.read()98 99def resize_image_to_bucket(image: Union[Image.Image, np.ndarray], bucket_reso: Tuple[int, int]) -> np.ndarray:100 """101 Resize the image to the bucket resolution.102 """103 if isinstance(image, Image.Image):104 image = np.array(image)105 elif not isinstance(image, np.ndarray):106 raise ValueError("Image must be a PIL Image or NumPy array")107 image_height, image_width = image.shape[:2]108 if bucket_reso == (image_width, image_height):109 return image110 bucket_width, bucket_height = bucket_reso111 scale_width = bucket_width / image_width112 scale_height = bucket_height / image_height113 scale = max(scale_width, scale_height)114 image_width = int(image_width * scale + 0.5)115 image_height = int(image_height * scale + 0.5)116 if scale > 1:117 image = Image.fromarray(image)118 image = image.resize((image_width, image_height), Image.LANCZOS)119 image = np.array(image)120 else:121 image = cv2.resize(image, (image_width, image_height), interpolation=cv2.INTER_AREA)122 # crop the image to the bucket resolution123 crop_left = (image_width - bucket_width) // 2124 crop_top = (image_height - bucket_height) // 2125 image = image[crop_top:crop_top + bucket_height, crop_left:crop_left + bucket_width]126 return image127 128# 129# @torch.inference_mode()130@spaces.GPU(duration=120)131def generate_video(prompt: str, frame1: Image.Image, frame2: Image.Image, resolution: str, guidance_scale: float, num_frames: int, num_inference_steps: int) -> bytes:132 # Debugging print statements133 print(f"Frame 1 Type: {type(frame1)}")134 print(f"Frame 2 Type: {type(frame2)}")135 print(f"Resolution: {resolution}")136 # Parse resolution137 width, height = map(int, resolution.split('x'))138 # Load and preprocess frames139 cond_frame1 = np.array(frame1)140 cond_frame2 = np.array(frame2)141 cond_frame1 = resize_image_to_bucket(cond_frame1, bucket_reso=(width, height))142 cond_frame2 = resize_image_to_bucket(cond_frame2, bucket_reso=(width, height))143 cond_video = np.zeros(shape=(num_frames, height, width, 3))144 cond_video[0], cond_video[-1] = cond_frame1, cond_frame2145 cond_video = torch.from_numpy(cond_video.copy()).permute(0, 3, 1, 2)146 cond_video = torch.stack([video_transforms(x) for x in cond_video], dim=0).unsqueeze(0)147 with torch.no_grad():148 image_or_video = cond_video.to(device="cuda", dtype=pipe.dtype)149 image_or_video = image_or_video.permute(0, 2, 1, 3, 4).contiguous() # [B, F, C, H, W] -> [B, C, F, H, W]150 cond_latents = pipe.vae.encode(image_or_video).latent_dist.sample()151 cond_latents = cond_latents * pipe.vae.config.scaling_factor152 cond_latents = cond_latents.to(dtype=pipe.dtype)153 assert not torch.any(torch.isnan(cond_latents))154 # Generate video155 video = call_pipe(156 pipe,157 prompt=prompt,158 num_frames=num_frames,159 num_inference_steps=num_inference_steps,160 image_latents=cond_latents,161 width=width,162 height=height,163 guidance_scale=guidance_scale,164 generator=torch.Generator(device="cuda").manual_seed(0),165 ).frames[0]166 # Export to video167 # TO-DO: Implement alternate method168 video_path = "output.mp4"169 export_to_video(video, video_path, fps=24)170 torch.cuda.empty_cache()171 return video_path172 173@torch.inference_mode()174def call_pipe(175 pipe,176 prompt: Union[str, List[str]] = None,177 prompt_2: Union[str, List[str]] = None,178 height: int = 720,179 width: int = 1280,180 num_frames: int = 129,181 num_inference_steps: int = 50,182 sigmas: Optional[List[float]] = None,183 guidance_scale: float = 6.0,184 num_videos_per_prompt: Optional[int] = 1,185 generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,186 latents: Optional[torch.Tensor] = None,187 prompt_embeds: Optional[torch.Tensor] = None,188 pooled_prompt_embeds: Optional[torch.Tensor] = None,189 prompt_attention_mask: Optional[torch.Tensor] = None,190 output_type: Optional[str] = "pil",191 return_dict: bool = True,192 attention_kwargs: Optional[dict] = None,193 callback_on_step_end: Optional[Union[callable, PipelineCallback, MultiPipelineCallbacks]] = None,194 callback_on_step_end_tensor_inputs: Optional[List[str]] = None,195 prompt_template: Optional[dict] = DEFAULT_PROMPT_TEMPLATE,196 max_sequence_length: int = 256,197 image_latents: Optional[torch.Tensor] = None,198):199 if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):200 callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs201 # 1. Check inputs. Raise error if not correct202 pipe.check_inputs(203 prompt,204 prompt_2,205 height,206 width,207 prompt_embeds,208 callback_on_step_end_tensor_inputs,209 prompt_template,210 )211 pipe._guidance_scale = guidance_scale212 pipe._attention_kwargs = attention_kwargs213 pipe._current_timestep = None214 pipe._interrupt = False215 device = pipe._execution_device216 # 2. Define call parameters217 if prompt is not None and isinstance(prompt, str):218 batch_size = 1219 elif prompt is not None and isinstance(prompt, list):220 batch_size = len(prompt)221 else:222 batch_size = prompt_embeds.shape[0]223 # 3. Encode input prompt224 prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = pipe.encode_prompt(225 prompt=prompt,226 prompt_2=prompt_2,227 prompt_template=prompt_template,228 num_videos_per_prompt=num_videos_per_prompt,229 prompt_embeds=prompt_embeds,230 pooled_prompt_embeds=pooled_prompt_embeds,231 prompt_attention_mask=prompt_attention_mask,232 device=device,233 max_sequence_length=max_sequence_length,234 )235 transformer_dtype = pipe.transformer.dtype236 prompt_embeds = prompt_embeds.to(transformer_dtype)237 prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)238 if pooled_prompt_embeds is not None:239 pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype)240 # 4. Prepare timesteps241 sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas242 timesteps, num_inference_steps = retrieve_timesteps(243 pipe.scheduler,244 num_inference_steps,245 device,246 sigmas=sigmas,247 )248 # 5. Prepare latent variables249 num_channels_latents = pipe.transformer.config.in_channels250 num_latent_frames = (num_frames - 1) // pipe.vae_scale_factor_temporal + 1251 latents = pipe.prepare_latents(252 batch_size * num_videos_per_prompt,253 num_channels_latents,254 height,255 width,256 num_latent_frames,257 torch.float32,258 device,259 generator,260 latents,261 )262 # 6. Prepare guidance condition263 guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0264 # 7. Denoising loop265 num_warmup_steps = len(timesteps) - num_inference_steps * pipe.scheduler.order266 pipe._num_timesteps = len(timesteps)267 pipe.text_encoder.to("cpu")268 pipe.text_encoder_2.to("cpu") 269 torch.cuda.empty_cache()270 with pipe.progress_bar(total=num_inference_steps) as progress_bar:271 for i, t in enumerate(timesteps):272 if pipe.interrupt:273 continue274 pipe._current_timestep = t275 latent_model_input = latents.to(transformer_dtype)276 timestep = t.expand(latents.shape[0]).to(latents.dtype)277 noise_pred = pipe.transformer(278 hidden_states=torch.cat([latent_model_input, image_latents], dim=1),279 timestep=timestep,280 encoder_hidden_states=prompt_embeds,281 encoder_attention_mask=prompt_attention_mask,282 pooled_projections=pooled_prompt_embeds,283 guidance=guidance,284 attention_kwargs=attention_kwargs,285 return_dict=False,286 )[0]287 # compute the previous noisy sample x_t -> x_t-1288 latents = pipe.scheduler.step(noise_pred, t, latents, return_dict=False)[0]289 if callback_on_step_end is not None:290 callback_kwargs = {}291 for k in callback_on_step_end_tensor_inputs:292 callback_kwargs[k] = locals()[k]293 callback_outputs = callback_on_step_end(pipe, i, t, callback_kwargs)294 latents = callback_outputs.pop("latents", latents)295 prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)296 # call the callback, if provided297 if i < len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipe.scheduler.order == 0):298 progress_bar.update()299 pipe._current_timestep = None300 if not output_type == "latent":301 latents = latents.to(pipe.vae.dtype) / pipe.vae.config.scaling_factor302 video = pipe.vae.decode(latents, return_dict=False)[0]303 video = pipe.video_processor.postprocess_video(video, output_type=output_type)304 else:305 video = latents306 # Offload all models307 pipe.maybe_free_model_hooks()308 if not return_dict:309 return (video,)310 return HunyuanVideoPipelineOutput(frames=video)311 312def main():313 # Define the interface inputs314 315 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:316 with gr.Group():317 gr.Markdown("""# HunyuanVideo Keyframe Control Lora for Video Generation318 **Generate videos using the HunyuanVideo model with a prompt and two (or more) frames as conditions. Gradio / HF Spaces implementation demo.**319 ---320 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.321 322 * 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.. 323 324 I will fill out a request for GPU allocation for the demo with HF soon.325 326 """)327 328 329 330 with gr.Row():331 with gr.Column(scale=5):332 with gr.Row():333 prompt_textbox = gr.Textbox(label="Prompt", value="a subject ...", scale=2)334 resolution = gr.Dropdown(335 label="Resolution",336 choices=["720x1280", "544x960", "1280x720", "960x544", "720x720"],337 value="544x960"338 )339 frame1 = gr.Image(label="Frame 1", type="pil")340 frame2 = gr.Image(label="Frame 2", type="pil")341 num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=30)342 guidance_scale = gr.Slider(minimum=0.1, maximum=20, step=0.1, label="Guidance Scale", value=6.0)343 num_frames = gr.Slider(minimum=1, maximum=129, step=1, label="Number of Frames", value=49)344 generate_button = gr.Button("Generate Video")345 with gr.Column(scale=3):346 outputs = gr.Video(label="Generated Video")347 with gr.Accordion(label="Examples"):348 markdown_content = read_markdown_file("examples.md")349 gr.Markdown(markdown_content, sanitize_html=False)350 with gr.Accordion():351 gr.Markdown("""352 353 ## HunyuanVideo Keyframe Control Lora is an adapter for HunyuanVideo T2V model for keyframe-based video generation.354 ---355 **Our architecture builds upon existing models, introducing key enhancements to optimize keyframe-based video generation**:356 357 * 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.358 * 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.359 * 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.360 361 ## Recommended Settings362 1. The model works best on human subjects. Single subject images work slightly better.363 2. It is recommended to use the following image generation resolutions `720x1280`, `544x960`, `1280x720`, `960x544`.364 3. It is recommended to set frames from 33 upto 97. Can go upto 121 frames as well (but not tested much).365 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.366 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.367 368 ## 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. 369 370 """, sanitize_html=False, elem_id="md_footer", container=True)371 372 generate_button.click(generate_video, inputs=[prompt_textbox, frame1, frame2, resolution, guidance_scale, num_frames, num_inference_steps], outputs=outputs)373 374 demo.launch(show_error=True)375 376if __name__ == "__main__":377 main()