# =============================================================================
# LTX-2 SCAIL-2 Character Animation (LoRA + mask channels) Training Configuration
# =============================================================================
#
# Trains SCAIL-2-style end-to-end character animation: a driving video latent is
# concatenated into the token sequence with a RoPE width offset (ΔW), and
# in-context mask channels (1 environment switch + K binding slots) are attached
# to the driving tokens to route motion per character.
#
# This combines LoRA on the attention/FFN blocks with an *unfrozen* widened
# patchify_proj (its new mask-channel input columns cannot be reached by LoRA and
# are trained directly). Set `model.mask_conditioning_channels` to the channel
# count = temporal_factor * (K + 1) = 8 * (6 + 1) = 56 for the default K=6.
#
# Dataset structure:
#   preprocessed_data_root/
#   ├── latents/           # Target video latents (what the model generates)
#   ├── conditions/        # Text embeddings for each video
#   ├── driving_latents/   # Driving video latents (same F/H/W as target)
#   └── char_masks/        # Semantic masks per sample, "mask" = [K+1, F_pix, H_pix, W_pix]
#                          #   channel 0 = environment switch, 1..K = character binding slots
#
# =============================================================================

model:
  model_path: "path/to/ltx-2-model.safetensors"
  text_encoder_path: "path/to/gemma-text-encoder"
  training_mode: "lora"
  # Widen the video patchify_proj by 56 zero-init input columns (8 * (K+1), K=6).
  # In LoRA mode the trainer additionally unfreezes patchify_proj so these train.
  mask_conditioning_channels: 56
  load_checkpoint: null

lora:
  rank: 32
  alpha: 32
  dropout: 0.0
  target_modules:
    - "attn1.to_k"
    - "attn1.to_q"
    - "attn1.to_v"
    - "attn1.to_out.0"
    - "attn2.to_k"
    - "attn2.to_q"
    - "attn2.to_v"
    - "attn2.to_out.0"
    - "ff.net.0.proj"
    - "ff.net.2"

training_strategy:
  name: "flexible"
  video:
    is_generated: true
    latents_dir: "latents"
    conditions:
      # SCAIL-2 driving conditioning: concatenate driving latents with a RoPE width
      # offset so they stay spatially detached from the target tokens.
      - type: driving
        latents_dir: "driving_latents"
        mode: "animation"        # "animation" | "replacement"
        width_offset: null        # null = target pixel width (driving sits just to the right)
        probability: 1.0
      # In-context mask channels attached to the driving tokens (target stays zero-mask).
      - type: mask_channels
        mask_dir: "char_masks"
        num_slots: 6              # K binding slots; channels = 8 * (K + 1) = 56 (match model.mask_conditioning_channels)

optimization:
  learning_rate: 2e-4
  steps: 3000
  batch_size: 1
  gradient_accumulation_steps: 1
  max_grad_norm: 1.0
  optimizer_type: "adamw"
  scheduler_type: "linear"
  scheduler_params: { }
  enable_gradient_checkpointing: true

acceleration:
  mixed_precision_mode: "bf16"
  quantization: null
  load_text_encoder_in_8bit: false
  offload_optimizer_during_validation: false

data:
  preprocessed_data_root: "/path/to/preprocessed/data"
  num_dataloader_workers: 2

validation:
  # NOTE: SCAIL driving + mask-channel validation conditions are not yet wired into
  # the validation runner (Phase 3 training path only). Keep validation minimal /
  # disabled until the inference pipeline (Phase 4) lands.
  samples:
    - prompt: >-
        A person performing an energetic dance routine, matching the motion of the driving
        performer, with crisp footwork and expressive arm movements in a bright studio.
      conditions: []
  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
  video_dims: [ 512, 512, 81 ]
  frame_rate: 25.0
  seed: 42
  inference_steps: 30
  interval: null   # disabled: driving/mask validation lands in Phase 4
  guidance_scale: 4.0
  stg_scale: 1.0
  stg_blocks: [29]
  stg_mode: "stg_v"
  generate_audio: false
  skip_initial_validation: true

checkpoints:
  interval: 250
  keep_last_n: 3
  precision: "bfloat16"

flow_matching:
  timestep_sampling_mode: "shifted_logit_normal"
  timestep_sampling_params: { }

hub:
  push_to_hub: false
  hub_model_id: null

wandb:
  enabled: false
  project: "ltx-2-trainer"
  entity: null
  tags: [ "ltx2", "scail-2", "character-animation" ]
  log_validation_videos: true

seed: 42
output_dir: "outputs/scail_animation_lora"