| 1 | import logging |
| 2 | from collections.abc import Iterator |
| 3 | |
| 4 | import torch |
| 5 | from einops import rearrange |
| 6 | from safetensors import safe_open |
| 7 | |
| 8 | from ltx_core.components.diffusion_steps import EulerDiffusionStep |
| 9 | from ltx_core.components.noisers import GaussianNoiser |
| 10 | from ltx_core.components.protocols import DiffusionStepProtocol |
| 11 | from ltx_core.conditioning import ( |
| 12 | ConditioningItem, |
| 13 | ConditioningItemAttentionStrengthWrapper, |
| 14 | VideoConditionByReferenceLatent, |
| 15 | ) |
| 16 | from ltx_core.loader import LoraPathStrengthAndSDOps |
| 17 | from ltx_core.model.audio_vae import decode_audio as vae_decode_audio |
| 18 | from ltx_core.model.upsampler import upsample_video |
| 19 | from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number |
| 20 | from ltx_core.model.video_vae import decode_video as vae_decode_video |
| 21 | from ltx_core.quantization import QuantizationPolicy |
| 22 | from ltx_core.types import Audio, LatentState, VideoLatentShape, VideoPixelShape |
| 23 | from ltx_pipelines.utils import ( |
| 24 | ModelLedger, |
| 25 | assert_resolution, |
| 26 | cleanup_memory, |
| 27 | combined_image_conditionings, |
| 28 | denoise_audio_video, |
| 29 | encode_prompts, |
| 30 | euler_denoising_loop, |
| 31 | get_device, |
| 32 | simple_denoising_func, |
| 33 | ) |
| 34 | from ltx_pipelines.utils.args import ( |
| 35 | ImageConditioningInput, |
| 36 | VideoConditioningAction, |
| 37 | VideoMaskConditioningAction, |
| 38 | default_2_stage_distilled_arg_parser, |
| 39 | detect_checkpoint_path, |
| 40 | ) |
| 41 | from ltx_pipelines.utils.constants import ( |
| 42 | DISTILLED_SIGMA_VALUES, |
| 43 | STAGE_2_DISTILLED_SIGMA_VALUES, |
| 44 | detect_params, |
| 45 | ) |
| 46 | from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning |
| 47 | from ltx_pipelines.utils.types import PipelineComponents |
| 48 | |
| 49 | device = get_device() |
| 50 | |
| 51 | |
| 52 | class ICLoraPipeline: |
| 53 | """ |
| 54 | Two-stage video generation pipeline with In-Context (IC) LoRA support. |
| 55 | Allows conditioning the generated video on control signals such as depth maps, |
| 56 | human pose, or image edges via the video_conditioning parameter. |
| 57 | The specific IC-LoRA model should be provided via the loras parameter. |
| 58 | Stage 1 generates video at half of the target resolution, then Stage 2 upsamples |
| 59 | by 2x and refines with additional denoising steps for higher quality output. |
| 60 | Both stages use distilled models for efficiency. |
| 61 | """ |
| 62 | |
| 63 | def __init__( |
| 64 | self, |
| 65 | distilled_checkpoint_path: str, |
| 66 | spatial_upsampler_path: str, |
| 67 | gemma_root: str, |
| 68 | loras: list[LoraPathStrengthAndSDOps], |
| 69 | device: torch.device = device, |
| 70 | quantization: QuantizationPolicy | None = None, |
| 71 | ): |
| 72 | self.dtype = torch.bfloat16 |
| 73 | self.stage_1_model_ledger = ModelLedger( |
| 74 | dtype=self.dtype, |
| 75 | device=device, |
| 76 | checkpoint_path=distilled_checkpoint_path, |
| 77 | spatial_upsampler_path=spatial_upsampler_path, |
| 78 | gemma_root_path=gemma_root, |
| 79 | loras=loras, |
| 80 | quantization=quantization, |
| 81 | ) |
| 82 | self.stage_2_model_ledger = ModelLedger( |
| 83 | dtype=self.dtype, |
| 84 | device=device, |
| 85 | checkpoint_path=distilled_checkpoint_path, |
| 86 | spatial_upsampler_path=spatial_upsampler_path, |
| 87 | gemma_root_path=gemma_root, |
| 88 | loras=[], |
| 89 | quantization=quantization, |
| 90 | ) |
| 91 | self.pipeline_components = PipelineComponents( |
| 92 | dtype=self.dtype, |
| 93 | device=device, |
| 94 | ) |
| 95 | self.device = device |
| 96 | |
| 97 | # Read reference downscale factor from LoRA metadata. |
| 98 | # IC-LoRAs trained with low-resolution reference videos store this factor |
| 99 | # so inference can resize reference videos to match training conditions. |
| 100 | self.reference_downscale_factor = 1 |
| 101 | for lora in loras: |
| 102 | scale = _read_lora_reference_downscale_factor(lora.path) |
| 103 | if scale != 1: |
| 104 | if self.reference_downscale_factor not in (1, scale): |
| 105 | raise ValueError( |
| 106 | f"Conflicting reference_downscale_factor values in LoRAs: " |
| 107 | f"already have {self.reference_downscale_factor}, but {lora.path} " |
| 108 | f"specifies {scale}. Cannot combine LoRAs with different reference scales." |
| 109 | ) |
| 110 | self.reference_downscale_factor = scale |
| 111 | |
| 112 | def __call__( # noqa: PLR0913 |
| 113 | self, |
| 114 | prompt: str, |
| 115 | seed: int, |
| 116 | height: int, |
| 117 | width: int, |
| 118 | num_frames: int, |
| 119 | frame_rate: float, |
| 120 | images: list[ImageConditioningInput], |
| 121 | video_conditioning: list[tuple[str, float]], |
| 122 | enhance_prompt: bool = False, |
| 123 | tiling_config: TilingConfig | None = None, |
| 124 | conditioning_attention_strength: float = 1.0, |
| 125 | skip_stage_2: bool = False, |
| 126 | conditioning_attention_mask: torch.Tensor | None = None, |
| 127 | ) -> tuple[Iterator[torch.Tensor], Audio]: |
| 128 | """ |
| 129 | Generate video with IC-LoRA conditioning. |
| 130 | Args: |
| 131 | prompt: Text prompt for video generation. |
| 132 | seed: Random seed for reproducibility. |
| 133 | height: Output video height in pixels (must be divisible by 64). |
| 134 | width: Output video width in pixels (must be divisible by 64). |
| 135 | num_frames: Number of frames to generate. |
| 136 | frame_rate: Output video frame rate. |
| 137 | images: List of (path, frame_idx, strength) tuples for image conditioning. |
| 138 | video_conditioning: List of (path, strength) tuples for IC-LoRA video conditioning. |
| 139 | enhance_prompt: Whether to enhance the prompt using the text encoder. |
| 140 | tiling_config: Optional tiling configuration for VAE decoding. |
| 141 | conditioning_attention_strength: Scale factor for IC-LoRA conditioning attention. |
| 142 | Controls how strongly the conditioning video influences the output. |
| 143 | 0.0 = ignore conditioning, 1.0 = full conditioning influence. Default 1.0. |
| 144 | When conditioning_attention_mask is provided, the mask is multiplied by |
| 145 | this strength before being passed to the conditioning items. |
| 146 | skip_stage_2: If True, skip Stage 2 upsampling and refinement. Output will be |
| 147 | at half resolution (height//2, width//2). Default is False. |
| 148 | conditioning_attention_mask: Optional pixel-space attention mask with the same |
| 149 | spatial-temporal dimensions as the input reference video. Shape should be |
| 150 | (B, 1, F, H, W) or (1, 1, F, H, W) where F, H, W match the reference |
| 151 | video's pixel dimensions. Values in [0, 1]. |
| 152 | The mask is downsampled to latent space using VAE scale factors (with |
| 153 | causal temporal handling for the first frame), then multiplied by |
| 154 | conditioning_attention_strength. |
| 155 | When None (default): scalar conditioning_attention_strength is used |
| 156 | directly. |
| 157 | Returns: |
| 158 | Tuple of (video_iterator, audio_tensor). |
| 159 | """ |
| 160 | assert_resolution(height=height, width=width, is_two_stage=True) |
| 161 | if not (0.0 <= conditioning_attention_strength <= 1.0): |
| 162 | raise ValueError( |
| 163 | f"conditioning_attention_strength must be in [0.0, 1.0], got {conditioning_attention_strength}" |
| 164 | ) |
| 165 | |
| 166 | generator = torch.Generator(device=self.device).manual_seed(seed) |
| 167 | noiser = GaussianNoiser(generator=generator) |
| 168 | stepper = EulerDiffusionStep() |
| 169 | dtype = torch.bfloat16 |
| 170 | |
| 171 | (ctx_p,) = encode_prompts( |
| 172 | [prompt], |
| 173 | self.stage_1_model_ledger, |
| 174 | enhance_first_prompt=enhance_prompt, |
| 175 | enhance_prompt_image=images[0][0] if len(images) > 0 else None, |
| 176 | enhance_prompt_seed=seed, |
| 177 | ) |
| 178 | video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding |
| 179 | |
| 180 | # Stage 1: Initial low resolution video generation. |
| 181 | stage_1_output_shape = VideoPixelShape( |
| 182 | batch=1, |
| 183 | frames=num_frames, |
| 184 | width=width // 2, |
| 185 | height=height // 2, |
| 186 | fps=frame_rate, |
| 187 | ) |
| 188 | |
| 189 | # Encode conditionings before loading transformer to reduce peak VRAM |
| 190 | video_encoder = self.stage_1_model_ledger.video_encoder() |
| 191 | stage_1_conditionings = self._create_conditionings( |
| 192 | images=images, |
| 193 | video_conditioning=video_conditioning, |
| 194 | height=stage_1_output_shape.height, |
| 195 | width=stage_1_output_shape.width, |
| 196 | video_encoder=video_encoder, |
| 197 | num_frames=num_frames, |
| 198 | conditioning_attention_strength=conditioning_attention_strength, |
| 199 | conditioning_attention_mask=conditioning_attention_mask, |
| 200 | ) |
| 201 | |
| 202 | transformer = self.stage_1_model_ledger.transformer() |
| 203 | stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device) |
| 204 | |
| 205 | def first_stage_denoising_loop( |
| 206 | sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol |
| 207 | ) -> tuple[LatentState, LatentState]: |
| 208 | return euler_denoising_loop( |
| 209 | sigmas=sigmas, |
| 210 | video_state=video_state, |
| 211 | audio_state=audio_state, |
| 212 | stepper=stepper, |
| 213 | denoise_fn=simple_denoising_func( |
| 214 | video_context=video_context, |
| 215 | audio_context=audio_context, |
| 216 | transformer=transformer, # noqa: F821 |
| 217 | ), |
| 218 | ) |
| 219 | |
| 220 | video_state, audio_state = denoise_audio_video( |
| 221 | output_shape=stage_1_output_shape, |
| 222 | conditionings=stage_1_conditionings, |
| 223 | noiser=noiser, |
| 224 | sigmas=stage_1_sigmas, |
| 225 | stepper=stepper, |
| 226 | denoising_loop_fn=first_stage_denoising_loop, |
| 227 | components=self.pipeline_components, |
| 228 | dtype=dtype, |
| 229 | device=self.device, |
| 230 | ) |
| 231 | |
| 232 | torch.cuda.synchronize() |
| 233 | del transformer |
| 234 | cleanup_memory() |
| 235 | |
| 236 | if skip_stage_2: |
| 237 | # Skip Stage 2: Decode directly from Stage 1 output at half resolution |
| 238 | logging.info("[IC-LoRA] Skipping Stage 2 (--skip-stage-2 enabled)") |
| 239 | decoded_video = vae_decode_video( |
| 240 | video_state.latent, self.stage_1_model_ledger.video_decoder(), tiling_config, generator |
| 241 | ) |
| 242 | decoded_audio = vae_decode_audio( |
| 243 | audio_state.latent, self.stage_1_model_ledger.audio_decoder(), self.stage_1_model_ledger.vocoder() |
| 244 | ) |
| 245 | del video_encoder |
| 246 | cleanup_memory() |
| 247 | return decoded_video, decoded_audio |
| 248 | |
| 249 | # Stage 2: Upsample and refine the video at higher resolution with distilled LORA. |
| 250 | upscaled_video_latent = upsample_video( |
| 251 | latent=video_state.latent[:1], |
| 252 | video_encoder=video_encoder, |
| 253 | upsampler=self.stage_2_model_ledger.spatial_upsampler(), |
| 254 | ) |
| 255 | |
| 256 | torch.cuda.synchronize() |
| 257 | cleanup_memory() |
| 258 | |
| 259 | transformer = self.stage_2_model_ledger.transformer() |
| 260 | distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device) |
| 261 | |
| 262 | def second_stage_denoising_loop( |
| 263 | sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol |
| 264 | ) -> tuple[LatentState, LatentState]: |
| 265 | return euler_denoising_loop( |
| 266 | sigmas=sigmas, |
| 267 | video_state=video_state, |
| 268 | audio_state=audio_state, |
| 269 | stepper=stepper, |
| 270 | denoise_fn=simple_denoising_func( |
| 271 | video_context=video_context, |
| 272 | audio_context=audio_context, |
| 273 | transformer=transformer, # noqa: F821 |
| 274 | ), |
| 275 | ) |
| 276 | |
| 277 | stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) |
| 278 | stage_2_conditionings = combined_image_conditionings( |
| 279 | images=images, |
| 280 | height=stage_2_output_shape.height, |
| 281 | width=stage_2_output_shape.width, |
| 282 | video_encoder=video_encoder, |
| 283 | dtype=self.dtype, |
| 284 | device=self.device, |
| 285 | ) |
| 286 | |
| 287 | video_state, audio_state = denoise_audio_video( |
| 288 | output_shape=stage_2_output_shape, |
| 289 | conditionings=stage_2_conditionings, |
| 290 | noiser=noiser, |
| 291 | sigmas=distilled_sigmas, |
| 292 | stepper=stepper, |
| 293 | denoising_loop_fn=second_stage_denoising_loop, |
| 294 | components=self.pipeline_components, |
| 295 | dtype=dtype, |
| 296 | device=self.device, |
| 297 | noise_scale=distilled_sigmas[0], |
| 298 | initial_video_latent=upscaled_video_latent, |
| 299 | initial_audio_latent=audio_state.latent, |
| 300 | ) |
| 301 | |
| 302 | torch.cuda.synchronize() |
| 303 | del transformer |
| 304 | del video_encoder |
| 305 | cleanup_memory() |
| 306 | |
| 307 | decoded_video = vae_decode_video( |
| 308 | video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator |
| 309 | ) |
| 310 | decoded_audio = vae_decode_audio( |
| 311 | audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder() |
| 312 | ) |
| 313 | return decoded_video, decoded_audio |
| 314 | |
| 315 | def _create_conditionings( |
| 316 | self, |
| 317 | images: list[ImageConditioningInput], |
| 318 | video_conditioning: list[tuple[str, float]], |
| 319 | height: int, |
| 320 | width: int, |
| 321 | num_frames: int, |
| 322 | video_encoder: VideoEncoder, |
| 323 | conditioning_attention_strength: float = 1.0, |
| 324 | conditioning_attention_mask: torch.Tensor | None = None, |
| 325 | ) -> list[ConditioningItem]: |
| 326 | """ |
| 327 | Create conditioning items for video generation. |
| 328 | Args: |
| 329 | conditioning_attention_strength: Scalar attention weight in [0, 1]. |
| 330 | If conditioning_attention_mask is also provided, the downsampled mask |
| 331 | is multiplied by this strength. Otherwise this scalar is passed |
| 332 | directly as the attention mask. |
| 333 | conditioning_attention_mask: Optional pixel-space attention mask with shape |
| 334 | (B, 1, F_pixel, H_pixel, W_pixel) matching the reference video's |
| 335 | pixel dimensions. Downsampled to latent space with causal temporal |
| 336 | handling, then multiplied by conditioning_attention_strength. |
| 337 | Returns: |
| 338 | List of conditioning items. IC-LoRA conditionings are appended last. |
| 339 | """ |
| 340 | conditionings = combined_image_conditionings( |
| 341 | images=images, |
| 342 | height=height, |
| 343 | width=width, |
| 344 | video_encoder=video_encoder, |
| 345 | dtype=self.dtype, |
| 346 | device=self.device, |
| 347 | ) |
| 348 | |
| 349 | # Calculate scaled dimensions for reference video conditioning. |
| 350 | # IC-LoRAs trained with downscaled reference videos expect the same ratio at inference. |
| 351 | scale = self.reference_downscale_factor |
| 352 | if scale != 1 and (height % scale != 0 or width % scale != 0): |
| 353 | raise ValueError( |
| 354 | f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})" |
| 355 | ) |
| 356 | ref_height = height // scale |
| 357 | ref_width = width // scale |
| 358 | |
| 359 | for video_path, strength in video_conditioning: |
| 360 | # Load video at scaled-down resolution (if scale > 1) |
| 361 | video = load_video_conditioning( |
| 362 | video_path=video_path, |
| 363 | height=ref_height, |
| 364 | width=ref_width, |
| 365 | frame_cap=num_frames, |
| 366 | dtype=self.dtype, |
| 367 | device=self.device, |
| 368 | ) |
| 369 | encoded_video = video_encoder(video) |
| 370 | reference_video_shape = VideoLatentShape.from_torch_shape(encoded_video.shape) |
| 371 | |
| 372 | # Build attention_mask for ConditioningItemAttentionStrengthWrapper |
| 373 | if conditioning_attention_mask is not None: |
| 374 | # Downsample pixel-space mask to latent space, then scale by strength |
| 375 | latent_mask = self._downsample_mask_to_latent( |
| 376 | mask=conditioning_attention_mask, |
| 377 | target_latent_shape=reference_video_shape, |
| 378 | ) |
| 379 | attn_mask = latent_mask * conditioning_attention_strength |
| 380 | elif conditioning_attention_strength < 1.0: |
| 381 | # Use scalar strength only |
| 382 | attn_mask = conditioning_attention_strength |
| 383 | else: |
| 384 | attn_mask = None |
| 385 | |
| 386 | cond = VideoConditionByReferenceLatent( |
| 387 | latent=encoded_video, |
| 388 | downscale_factor=scale, |
| 389 | strength=strength, |
| 390 | ) |
| 391 | if attn_mask is not None: |
| 392 | cond = ConditioningItemAttentionStrengthWrapper(cond, attention_mask=attn_mask) |
| 393 | conditionings.append(cond) |
| 394 | |
| 395 | if video_conditioning: |
| 396 | logging.info(f"[IC-LoRA] Added {len(video_conditioning)} video conditioning(s)") |
| 397 | |
| 398 | return conditionings |
| 399 | |
| 400 | @staticmethod |
| 401 | def _downsample_mask_to_latent( |
| 402 | mask: torch.Tensor, |
| 403 | target_latent_shape: VideoLatentShape, |
| 404 | ) -> torch.Tensor: |
| 405 | """ |
| 406 | Downsample a pixel-space mask to latent space using VAE scale factors. |
| 407 | Handles causal temporal downsampling: the first frame is kept separately |
| 408 | (temporal scale factor = 1 for the first frame), while the remaining |
| 409 | frames are downsampled by the VAE's temporal scale factor. |
| 410 | Args: |
| 411 | mask: Pixel-space mask of shape (B, 1, F_pixel, H_pixel, W_pixel). |
| 412 | Values in [0, 1]. |
| 413 | target_latent_shape: Expected latent shape after VAE encoding. |
| 414 | Used to determine the target (F_latent, H_latent, W_latent). |
| 415 | Returns: |
| 416 | Flattened latent-space mask of shape (B, F_lat * H_lat * W_lat), |
| 417 | matching the patchifier's token ordering (f, h, w). |
| 418 | """ |
| 419 | b = mask.shape[0] |
| 420 | f_lat = target_latent_shape.frames |
| 421 | h_lat = target_latent_shape.height |
| 422 | w_lat = target_latent_shape.width |
| 423 | |
| 424 | # Step 1: Spatial downsampling (area interpolation per frame) |
| 425 | f_pix = mask.shape[2] |
| 426 | spatial_down = torch.nn.functional.interpolate( |
| 427 | rearrange(mask, "b 1 f h w -> (b f) 1 h w"), |
| 428 | size=(h_lat, w_lat), |
| 429 | mode="area", |
| 430 | ) |
| 431 | spatial_down = rearrange(spatial_down, "(b f) 1 h w -> b 1 f h w", b=b) |
| 432 | |
| 433 | # Step 2: Causal temporal downsampling |
| 434 | # First frame: kept as-is (causal VAE encodes first frame independently) |
| 435 | first_frame = spatial_down[:, :, :1, :, :] # (B, 1, 1, H_lat, W_lat) |
| 436 | |
| 437 | if f_pix > 1 and f_lat > 1: |
| 438 | # Remaining frames: downsample by temporal factor via group-mean |
| 439 | t = (f_pix - 1) // (f_lat - 1) # temporal downscale factor |
| 440 | assert (f_pix - 1) % (f_lat - 1) == 0, ( |
| 441 | f"Pixel frames ({f_pix}) not compatible with latent frames ({f_lat}): " |
| 442 | f"(f_pix - 1) must be divisible by (f_lat - 1)" |
| 443 | ) |
| 444 | rest = rearrange(spatial_down[:, :, 1:, :, :], "b 1 (f t) h w -> b 1 f t h w", t=t) |
| 445 | rest = rest.mean(dim=3) # (B, 1, F_lat-1, H_lat, W_lat) |
| 446 | latent_mask = torch.cat([first_frame, rest], dim=2) # (B, 1, F_lat, H_lat, W_lat) |
| 447 | else: |
| 448 | latent_mask = first_frame |
| 449 | |
| 450 | # Flatten to (B, F_lat * H_lat * W_lat) matching patchifier token order (f, h, w) |
| 451 | return rearrange(latent_mask, "b 1 f h w -> b (f h w)") |
| 452 | |
| 453 | |
| 454 | @torch.inference_mode() |
| 455 | def main() -> None: |
| 456 | logging.getLogger().setLevel(logging.INFO) |
| 457 | checkpoint_path = detect_checkpoint_path(distilled=True) |
| 458 | params = detect_params(checkpoint_path) |
| 459 | parser = default_2_stage_distilled_arg_parser(params=params) |
| 460 | parser.add_argument( |
| 461 | "--video-conditioning", |
| 462 | action=VideoConditioningAction, |
| 463 | nargs=2, |
| 464 | metavar=("PATH", "STRENGTH"), |
| 465 | required=True, |
| 466 | ) |
| 467 | parser.add_argument( |
| 468 | "--conditioning-attention-mask", |
| 469 | action=VideoMaskConditioningAction, |
| 470 | nargs=2, |
| 471 | metavar=("MASK_PATH", "STRENGTH"), |
| 472 | default=None, |
| 473 | help=( |
| 474 | "Optional spatial attention mask: path to a grayscale mask video and " |
| 475 | "attention strength. The mask video pixel values in [0,1] control " |
| 476 | "per-region conditioning attention strength. The strength scalar is " |
| 477 | "multiplied with the spatial mask. " |
| 478 | "0.0 = ignore IC-LoRA conditioning, 1.0 = full conditioning influence. " |
| 479 | "When not provided, full conditioning strength (1.0) is used. " |
| 480 | "Example: --conditioning-attention-mask path/to/mask.mp4 0.5" |
| 481 | ), |
| 482 | ) |
| 483 | parser.add_argument( |
| 484 | "--skip-stage-2", |
| 485 | action="store_true", |
| 486 | help=( |
| 487 | "Skip Stage 2 upsampling and refinement. Output will be at half resolution " |
| 488 | "(height//2, width//2). Useful for faster iteration or when GPU memory is limited." |
| 489 | ), |
| 490 | ) |
| 491 | args = parser.parse_args() |
| 492 | |
| 493 | # Load mask video if provided via --conditioning-attention-mask |
| 494 | conditioning_attention_mask = None |
| 495 | conditioning_attention_strength = 1.0 |
| 496 | if args.conditioning_attention_mask is not None: |
| 497 | mask_path, mask_strength = args.conditioning_attention_mask |
| 498 | conditioning_attention_strength = mask_strength |
| 499 | conditioning_attention_mask = _load_mask_video( |
| 500 | mask_path=mask_path, |
| 501 | height=args.height // 2, # Stage 1 operates at half resolution |
| 502 | width=args.width // 2, |
| 503 | num_frames=args.num_frames, |
| 504 | ) |
| 505 | |
| 506 | pipeline = ICLoraPipeline( |
| 507 | distilled_checkpoint_path=args.distilled_checkpoint_path, |
| 508 | spatial_upsampler_path=args.spatial_upsampler_path, |
| 509 | gemma_root=args.gemma_root, |
| 510 | loras=tuple(args.lora) if args.lora else (), |
| 511 | quantization=args.quantization, |
| 512 | ) |
| 513 | tiling_config = TilingConfig.default() |
| 514 | video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config) |
| 515 | video, audio = pipeline( |
| 516 | prompt=args.prompt, |
| 517 | seed=args.seed, |
| 518 | height=args.height, |
| 519 | width=args.width, |
| 520 | num_frames=args.num_frames, |
| 521 | frame_rate=args.frame_rate, |
| 522 | images=args.images, |
| 523 | video_conditioning=args.video_conditioning, |
| 524 | tiling_config=tiling_config, |
| 525 | conditioning_attention_strength=conditioning_attention_strength, |
| 526 | skip_stage_2=args.skip_stage_2, |
| 527 | conditioning_attention_mask=conditioning_attention_mask, |
| 528 | ) |
| 529 | |
| 530 | encode_video( |
| 531 | video=video, |
| 532 | fps=args.frame_rate, |
| 533 | audio=audio, |
| 534 | output_path=args.output_path, |
| 535 | video_chunks_number=video_chunks_number, |
| 536 | ) |
| 537 | |
| 538 | |
| 539 | def _load_mask_video( |
| 540 | mask_path: str, |
| 541 | height: int, |
| 542 | width: int, |
| 543 | num_frames: int, |
| 544 | ) -> torch.Tensor: |
| 545 | """Load a mask video and return a pixel-space tensor of shape (1, 1, F, H, W). |
| 546 | The mask video is loaded, resized to (height, width), converted to |
| 547 | grayscale, and normalised to [0, 1]. |
| 548 | Args: |
| 549 | mask_path: Path to the mask video file. |
| 550 | height: Target height in pixels. |
| 551 | width: Target width in pixels. |
| 552 | num_frames: Maximum number of frames to load. |
| 553 | Returns: |
| 554 | Tensor of shape ``(1, 1, F, H, W)`` with values in ``[0, 1]``. |
| 555 | """ |
| 556 | mask_video = load_video_conditioning( |
| 557 | video_path=mask_path, |
| 558 | height=height, |
| 559 | width=width, |
| 560 | frame_cap=num_frames, |
| 561 | dtype=torch.bfloat16, |
| 562 | device=device, |
| 563 | ) |
| 564 | # mask_video shape: (1, C, F, H, W) — take mean over channels for grayscale |
| 565 | mask = mask_video.mean(dim=1, keepdim=True) # (1, 1, F, H, W) |
| 566 | # Normalise to [0, 1] — load_video_conditioning applies normalize_latent, |
| 567 | # so undo that: values are in [-1, 1], remap to [0, 1] |
| 568 | mask = (mask + 1.0) / 2.0 |
| 569 | return mask.clamp(0.0, 1.0) |
| 570 | |
| 571 | |
| 572 | def _read_lora_reference_downscale_factor(lora_path: str) -> int: |
| 573 | """Read reference_downscale_factor from LoRA safetensors metadata. |
| 574 | Some IC-LoRA models are trained with reference videos at lower resolution than |
| 575 | the target output. This allows for more efficient training and can improve |
| 576 | generalization. The downscale factor indicates the ratio between target and |
| 577 | reference resolutions (e.g., factor=2 means reference is half the resolution). |
| 578 | Args: |
| 579 | lora_path: Path to the LoRA .safetensors file |
| 580 | Returns: |
| 581 | The reference downscale factor (1 if not specified in metadata, meaning |
| 582 | reference and target have the same resolution) |
| 583 | """ |
| 584 | try: |
| 585 | with safe_open(lora_path, framework="pt") as f: |
| 586 | metadata = f.metadata() or {} |
| 587 | return int(metadata.get("reference_downscale_factor", 1)) |
| 588 | except Exception as e: |
| 589 | logging.warning(f"Failed to read metadata from LoRA file '{lora_path}': {e}") |
| 590 | return 1 |
| 591 | |
| 592 | |
| 593 | if __name__ == "__main__": |
| 594 | main() |
| 595 |