Reducing cogload on the framework.

finetrainers/conditioning/conditioned_pipeline.py

@@ -13,6 +13,10 @@ from typing import Callable, Any
 import numpy as np
 from diffusers.utils import is_torch_xla_available
 from diffusers.pipelines.ltx.pipeline_output import LTXPipelineOutput
+from diffusers.utils import is_torch_xla_available
+from torch import torch
+from transformers import T5EncoderModel, T5TokenizerFast
+from typing import Any, Callable, Dict, List, Optional, Union
 
 
 if is_torch_xla_available():

finetrainers/ltx_video/__init__.py

@@ -1,2 +1,3 @@
 from .full_finetune import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG
+from .full_finetune_condition import LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG
 from .lora import LTX_VIDEO_T2V_LORA_CONFIG

finetrainers/ltx_video/full_finetune.py

@@ -1,4 +1,5 @@
 from diffusers import LTXPipeline
+from finetrainers.conditioning.conditioned_pipeline import LTXConditionedPipeline
 
 from .lora import (
     collate_fn_t2v,
@@ -6,29 +7,24 @@ from .lora import (
     initialize_pipeline,
     load_condition_models,
     load_diffusion_models,
-    load_conditioned_diffusion_models, # use this for conditioning
     load_latent_models,
     post_latent_preparation,
     prepare_conditions,
     prepare_latents,
     validation,
-    conditional_validation,
-    collate_fn_t2v_cond,
-    conditioned_forward_pass # use this for conditioning
 )
 
-
 # TODO(aryan): refactor into model specs for better re-use
 LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = {
     "pipeline_cls": LTXPipeline,
     "load_condition_models": load_condition_models,
     "load_latent_models": load_latent_models,
-    "load_diffusion_models": load_conditioned_diffusion_models,
+    "load_diffusion_models": load_diffusion_models,
     "initialize_pipeline": initialize_pipeline,
     "prepare_conditions": prepare_conditions,
     "prepare_latents": prepare_latents,
     "post_latent_preparation": post_latent_preparation,
-    "collate_fn": collate_fn_t2v_cond,
+    "collate_fn": collate_fn_t2v,
-    "forward_pass": conditioned_forward_pass,
+    "forward_pass": forward_pass,
-    "validation": conditional_validation,
+    "validation": validation,
 }

finetrainers/ltx_video/full_finetune_condition.py

@@ -0,0 +1,141 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+
+from diffusers import FlowMatchEulerDiscreteScheduler, LTXVideoTransformer3DModel
+from finetrainers.conditioning.LTXVideoConditionedTransformer3DModel import LTXVideoConditionedTransformer3DModel
+from finetrainers.conditioning.conditioned_pipeline import LTXConditionedPipeline
+from PIL import Image
+
+# Exisiting Pipeline
+from .lora import (
+    initialize_pipeline,
+    load_conditioned_diffusion_models, # use this for conditioning
+    load_latent_models,
+    post_latent_preparation,
+    prepare_conditions,
+    prepare_latents,
+    load_condition_models
+)
+
+def load_conditioned_diffusion_models(
+    model_id: str = "Lightricks/LTX-Video",
+    transformer_dtype: torch.dtype = torch.bfloat16,
+    revision: Optional[str] = None,
+    cache_dir: Optional[str] = None,
+    **kwargs,
+) -> Dict[str, nn.Module]:
+    transformer = LTXVideoConditionedTransformer3DModel.from_pretrained(
+        model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir
+    )
+    scheduler = FlowMatchEulerDiscreteScheduler()
+    return {"transformer": transformer, "scheduler": scheduler}
+
+
+def load_diffusion_models(
+    model_id: str = "Lightricks/LTX-Video",
+    transformer_dtype: torch.dtype = torch.bfloat16,
+    revision: Optional[str] = None,
+    cache_dir: Optional[str] = None,
+    **kwargs,
+) -> Dict[str, nn.Module]:
+    transformer = LTXVideoTransformer3DModel.from_pretrained(
+        model_id, subfolder="transformer", torch_dtype=transformer_dtype, revision=revision, cache_dir=cache_dir
+    )
+    scheduler = FlowMatchEulerDiscreteScheduler()
+    return {"transformer": transformer, "scheduler": scheduler}
+
+
+def collate_fn_t2v_cond(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+    return {
+        "prompts": [x["prompt"] for x in batch[0]],
+        "videos": torch.stack([x["video"] for x in batch[0]]),
+        "poses": torch.stack([x["pose"] for x in batch[0]]),
+        "img_refs": torch.stack([x["img_ref"] for x in batch[0]])
+    }
+
+def conditioned_forward_pass(
+    transformer: LTXVideoConditionedTransformer3DModel,
+    prompt_embeds: torch.Tensor,
+    prompt_attention_mask: torch.Tensor,
+    latents: torch.Tensor,
+    noisy_latents: torch.Tensor,
+    timesteps: torch.LongTensor,
+    num_frames: int,
+    height: int,
+    width: int,
+    noisy_latents_residual:torch.Tensor,
+    **kwargs,
+) -> torch.Tensor:
+    rope_interpolation_scale = [1 / 25, 32, 32]
+   
+    denoised_latents = transformer(
+        hidden_states=noisy_latents,
+        encoder_hidden_states=prompt_embeds,
+        timestep=timesteps,
+        encoder_attention_mask=prompt_attention_mask,
+        num_frames=num_frames,
+        height=height,
+        width=width,
+        rope_interpolation_scale=rope_interpolation_scale,
+        return_dict=False,
+        residual_x=noisy_latents_residual,
+    )[0]
+
+    return {"latents": denoised_latents}
+
+def conditional_validation(
+    pipeline: LTXConditionedPipeline,
+    prompt: str,
+    image: Optional[Image.Image] = None,
+    video: Optional[List[Image.Image]] = None,
+    pose_video: Optional[List[Image.Image]] = None,
+    image_ref_video: Optional[List[Image.Image]] = None,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_frames: Optional[int] = None,
+    frame_rate: int = 24,
+    num_videos_per_prompt: int = 1,
+    generator: Optional[torch.Generator] = None,
+    **kwargs,
+):
+    generation_kwargs = {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "num_frames": num_frames,
+        "frame_rate": frame_rate,
+        "num_videos_per_prompt": num_videos_per_prompt,
+        "generator": generator,
+        "return_dict": True,
+        "output_type": "pil",
+        "pose_video": pose_video,
+        "image_ref_video": image_ref_video,
+    }
+
+    # pose template noisey input [cat] img_ref + pose video 
+      
+    # pose template ref video latent + patchify video latent.
+    # img ref patchify video latent +
+    # add them together as input 
+    # residual x latent 
+
+    generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}
+    video = pipeline(**generation_kwargs).frames[0]
+    return [("video", video)]
+
+
+LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG = {
+    "pipeline_cls": LTXConditionedPipeline,
+    "load_condition_models": load_condition_models,
+    "load_latent_models": load_latent_models,
+    "load_diffusion_models": load_conditioned_diffusion_models,
+    "initialize_pipeline": initialize_pipeline,
+    "prepare_conditions": prepare_conditions,
+    "prepare_latents": prepare_latents,
+    "post_latent_preparation": post_latent_preparation,
+    "collate_fn": collate_fn_t2v_cond,
+    "forward_pass": conditioned_forward_pass,
+    "validation": conditional_validation,
+}

finetrainers/ltx_video/lora.py

@@ -5,13 +5,12 @@ import torch.nn as nn
 from accelerate.logging import get_logger
 from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel
 from finetrainers.conditioning.LTXVideoConditionedTransformer3DModel import LTXVideoConditionedTransformer3DModel
+from finetrainers.conditioning.conditioned_pipeline import LTXConditionedPipeline
 from PIL import Image
 from transformers import T5EncoderModel, T5Tokenizer
 
-
 logger = get_logger("finetrainers")  # pylint: disable=invalid-name
 
-
 def load_condition_models(
     model_id: str = "Lightricks/LTX-Video",
     text_encoder_dtype: torch.dtype = torch.bfloat16,
@@ -38,9 +37,9 @@ def load_latent_models(
     )
     return {"vae": vae}
 
-def load_conditioned_diffusion_models(
+def load_condition_models(
     model_id: str = "Lightricks/LTX-Video",
-    transformer_dtype: torch.dtype = torch.bfloat16,
+    text_encoder_dtype: torch.dtype = torch.bfloat16,
     revision: Optional[str] = None,
     cache_dir: Optional[str] = None,
     **kwargs,
@@ -84,7 +83,7 @@ def initialize_pipeline(
     enable_tiling: bool = False,
     enable_model_cpu_offload: bool = False,
     **kwargs,
-) -> LTXPipeline:
+) -> LTXConditionedPipeline:
     component_name_pairs = [
         ("tokenizer", tokenizer),
         ("text_encoder", text_encoder),
@@ -97,7 +96,7 @@ def initialize_pipeline(
         if component is not None:
             components[name] = component
 
-    pipe = LTXPipeline.from_pretrained(model_id, **components, revision=revision, cache_dir=cache_dir)
+    pipe = LTXConditionedPipeline.from_pretrained(model_id, **components, revision=revision, cache_dir=cache_dir)
     pipe.text_encoder = pipe.text_encoder.to(dtype=text_encoder_dtype)
     pipe.transformer = pipe.transformer.to(dtype=transformer_dtype)
     pipe.vae = pipe.vae.to(dtype=vae_dtype)
@@ -196,14 +195,6 @@ def post_latent_preparation(
     return {"latents": latents, "num_frames": num_frames, "height": height, "width": width}
 
 
-def collate_fn_t2v_cond(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
-    return {
-        "prompts": [x["prompt"] for x in batch[0]],
-        "videos": torch.stack([x["video"] for x in batch[0]]),
-        "poses": torch.stack([x["pose"] for x in batch[0]]),
-        "img_refs": torch.stack([x["img_ref"] for x in batch[0]])
-    }
-
 def collate_fn_t2v(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
     return {
         "prompts": [x["prompt"] for x in batch[0]],
@@ -211,35 +202,6 @@ def collate_fn_t2v(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torc
     }
 
 
-def conditioned_forward_pass(
-    transformer: LTXVideoConditionedTransformer3DModel,
-    prompt_embeds: torch.Tensor,
-    prompt_attention_mask: torch.Tensor,
-    latents: torch.Tensor,
-    noisy_latents: torch.Tensor,
-    timesteps: torch.LongTensor,
-    num_frames: int,
-    height: int,
-    width: int,
-    noisy_latents_residual:torch.Tensor,
-    **kwargs,
-) -> torch.Tensor:
-    rope_interpolation_scale = [1 / 25, 32, 32]
-   
-    denoised_latents = transformer(
-        hidden_states=noisy_latents,
-        encoder_hidden_states=prompt_embeds,
-        timestep=timesteps,
-        encoder_attention_mask=prompt_attention_mask,
-        num_frames=num_frames,
-        height=height,
-        width=width,
-        rope_interpolation_scale=rope_interpolation_scale,
-        return_dict=False,
-        residual_x=noisy_latents_residual,
-    )[0]
-
-    return {"latents": denoised_latents}
 
 def forward_pass(
     transformer: LTXVideoTransformer3DModel,
@@ -270,45 +232,7 @@ def forward_pass(
 
     return {"latents": denoised_latents}
 
-def conditional_validation(
-    pipeline: LTXPipeline,
-    prompt: str,
-    image: Optional[Image.Image] = None,
-    video: Optional[List[Image.Image]] = None,
-    pose_video: Optional[List[Image.Image]] = None,
-    image_ref_video: Optional[List[Image.Image]] = None,
-    height: Optional[int] = None,
-    width: Optional[int] = None,
-    num_frames: Optional[int] = None,
-    frame_rate: int = 24,
-    num_videos_per_prompt: int = 1,
-    generator: Optional[torch.Generator] = None,
-    **kwargs,
-):
-    generation_kwargs = {
-        "prompt": prompt,
-        "height": height,
-        "width": width,
-        "num_frames": num_frames,
-        "frame_rate": frame_rate,
-        "num_videos_per_prompt": num_videos_per_prompt,
-        "generator": generator,
-        "return_dict": True,
-        "output_type": "pil",
-        "pose_video": pose_video,
-        "image_ref_video": image_ref_video,
-    }
-    # pose template noisey input [cat] img_ref + pose video 
-      
-    # pose template ref video latent + patchify video latent.
-    # img ref patchify video latent +
-    # add them together as input 
-    # residual x latent 
 
-    generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}
-    video = pipeline(**generation_kwargs).frames[0]
-    
-    
 def validation(
     pipeline: LTXPipeline,
     prompt: str,

finetrainers/models.py

@@ -2,7 +2,7 @@ from typing import Any, Dict
 
 from .cogvideox import COGVIDEOX_T2V_FULL_FINETUNE_CONFIG, COGVIDEOX_T2V_LORA_CONFIG
 from .hunyuan_video import HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG, HUNYUAN_VIDEO_T2V_LORA_CONFIG
-from .ltx_video import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG, LTX_VIDEO_T2V_LORA_CONFIG
+from .ltx_video import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG,LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG, LTX_VIDEO_T2V_LORA_CONFIG
 
 
 SUPPORTED_MODEL_CONFIGS = {
@@ -12,7 +12,7 @@ SUPPORTED_MODEL_CONFIGS = {
     },
     "ltx_video": {
         "lora": LTX_VIDEO_T2V_LORA_CONFIG,
-        "full-finetune": LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG,
+        "full-finetune":LTX_VIDEO_COND_T2V_FULL_FINETUNE_CONFIG # LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG,
     },
     "cogvideox": {
         "lora": COGVIDEOX_T2V_LORA_CONFIG,