//! LTX Video inference engine — text-to-video generation.
//!
//! Architecture: T5-XXL text encoder → LTXVideoTransformer3DModel → 3D Causal VAE → APNG/GIF/WebP/MP4
//! Follows the same patterns as Flux2Engine (drop-and-reload, VRAM management, progress).

use std::sync::{Arc, Mutex};
use std::time::Instant;

use anyhow::{bail, Result};
use candle_core::{DType, Device, IndexOp, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::ltx_video::{
    sampling::{
        FlowMatchEulerDiscreteScheduler, FlowMatchEulerDiscreteSchedulerConfig, TimeShiftType,
    },
    transformer::{LtxVideoTransformer3DModel, LtxVideoTransformer3DModelConfig},
    vae::{AutoencoderKLLtxVideo, AutoencoderKLLtxVideoConfig},
};

use mold_core::{GenerateRequest, GenerateResponse, ModelPaths, OutputFormat, VideoData};

use crate::engine::{gpu_dtype, rand_seed, seeded_randn, LoadStrategy};
use crate::engine_base::EngineBase;
use crate::progress::{ProgressCallback, ProgressEvent};
use crate::shared_pool::SharedPool;

use super::{latent_upsampler::LatentUpsampler, video_enc};

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

/// Spatial compression ratio of the LTX Video VAE.
const VAE_SPATIAL_COMPRESSION: usize = 32;
/// Temporal compression ratio of the LTX Video VAE.
const VAE_TEMPORAL_COMPRESSION: usize = 8;
/// Latent channels in the VAE.
const LATENT_CHANNELS: usize = 128;
/// Patch sizes (both 1 for current LTX Video checkpoints).
const PATCH_SIZE: usize = 1;
const PATCH_SIZE_T: usize = 1;

const LTX_098_DISTILLED_FIRST_PASS_SIGMAS: &[f32] =
    &[1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250];
const LTX_098_DISTILLED_SECOND_PASS_SIGMAS: &[f32] = &[0.9094, 0.7250, 0.4219];
const LTX_098_DEV_FIRST_PASS_GUIDANCE_SCALE: &[f32] = &[1.0, 1.0, 6.0, 8.0, 6.0, 1.0, 1.0];
const LTX_098_DEV_FIRST_PASS_STG_SCALE: &[f32] = &[0.0, 0.0, 4.0, 4.0, 4.0, 2.0, 1.0];
const LTX_098_DEV_FIRST_PASS_RESCALING_SCALE: &[f32] = &[1.0, 1.0, 0.5, 0.5, 1.0, 1.0, 1.0];
const LTX_098_DEV_FIRST_PASS_GUIDANCE_TIMESTEPS: &[f32] =
    &[1.0, 0.996, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180];
const LTX_096_DEV_SKIP_BLOCKS: &[usize] = &[19];
const LTX_098_2B_DISTILLED_SKIP_BLOCKS: &[usize] = &[];
const LTX_098_13B_DISTILLED_SKIP_BLOCKS: &[usize] = &[42];
const LTX_098_13B_DEV_FIRST_PASS_SKIP_BLOCKS: &[&[usize]] = &[
    &[],
    &[11, 25, 35, 39],
    &[22, 35, 39],
    &[28],
    &[28],
    &[28],
    &[28],
];
const LTX_098_13B_DEV_SECOND_PASS_SKIP_BLOCKS: &[usize] = &[27];

fn is_official_ltx_transformer_checkpoint(path: &std::path::Path) -> bool {
    path.file_name()
        .and_then(|name| name.to_str())
        .is_some_and(|name| {
            name.ends_with(".safetensors")
                && name.starts_with("ltx")
                && !name.starts_with("diffusion_pytorch_model")
        })
}

fn remap_official_ltx_transformer_key(key: &str) -> String {
    let key = key
        .replace("proj_in", "patchify_proj")
        .replace("time_embed", "adaln_single")
        .replace("norm_q", "q_norm")
        .replace("norm_k", "k_norm");
    format!("model.diffusion_model.{key}")
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum LtxPipelineMode {
    Base,
    Multiscale,
}

#[derive(Clone, Debug)]
struct LtxGuidanceConfig {
    guidance_scale: Vec<f32>,
    stg_scale: Vec<f32>,
    rescaling_scale: Vec<f32>,
    guidance_timesteps: Option<Vec<f32>>,
    skip_block_list: Vec<Vec<usize>>,
    cfg_star_rescale: bool,
}

impl LtxGuidanceConfig {
    fn constant(
        guidance_scale: f32,
        stg_scale: f32,
        rescaling_scale: f32,
        skip_block_list: &[usize],
    ) -> Self {
        Self {
            guidance_scale: vec![guidance_scale],
            stg_scale: vec![stg_scale],
            rescaling_scale: vec![rescaling_scale],
            guidance_timesteps: None,
            skip_block_list: vec![skip_block_list.to_vec()],
            cfg_star_rescale: false,
        }
    }
}

#[derive(Clone, Debug)]
struct LtxPassConfig {
    num_inference_steps: u32,
    custom_sigmas: Option<Vec<f32>>,
    skip_initial_inference_steps: usize,
    skip_final_inference_steps: usize,
    guidance: LtxGuidanceConfig,
    tone_map_compression_ratio: f32,
}

#[derive(Clone, Debug)]
struct LtxMultiscaleConfig {
    downscale_factor: f32,
    first_pass: LtxPassConfig,
    second_pass: LtxPassConfig,
}

#[derive(Clone, Debug)]
struct LtxModelPreset {
    transformer_config: LtxVideoTransformer3DModelConfig,
    vae_config: AutoencoderKLLtxVideoConfig,
    scheduler_config: FlowMatchEulerDiscreteSchedulerConfig,
    base_pass: LtxPassConfig,
    decode_timestep: f32,
    decode_noise_scale: f32,
    mode: LtxPipelineMode,
    multiscale: Option<LtxMultiscaleConfig>,
}

impl LtxModelPreset {
    fn for_model(model_name: &str) -> Result<Self> {
        if model_name.contains("ltx-video-0.9.6-distilled") {
            Ok(Self {
                transformer_config: transformer_2b_config(),
                vae_config: improved_vae_config(),
                scheduler_config: scheduler_config(true),
                base_pass: LtxPassConfig {
                    num_inference_steps: 8,
                    custom_sigmas: None,
                    skip_initial_inference_steps: 0,
                    skip_final_inference_steps: 0,
                    guidance: LtxGuidanceConfig::constant(1.0, 0.0, 1.0, &[]),
                    tone_map_compression_ratio: 0.0,
                },
                decode_timestep: 0.05,
                decode_noise_scale: 0.025,
                mode: LtxPipelineMode::Base,
                multiscale: None,
            })
        } else if model_name.contains("ltx-video-0.9.6") {
            Ok(Self {
                transformer_config: transformer_2b_config(),
                vae_config: improved_vae_config(),
                scheduler_config: scheduler_config(false),
                base_pass: LtxPassConfig {
                    num_inference_steps: 40,
                    custom_sigmas: None,
                    skip_initial_inference_steps: 0,
                    skip_final_inference_steps: 0,
                    guidance: LtxGuidanceConfig::constant(3.0, 1.0, 0.7, LTX_096_DEV_SKIP_BLOCKS),
                    tone_map_compression_ratio: 0.0,
                },
                decode_timestep: 0.05,
                decode_noise_scale: 0.025,
                mode: LtxPipelineMode::Base,
                multiscale: None,
            })
        } else if model_name.contains("ltx-video-0.9.8-2b-distilled") {
            Ok(Self {
                transformer_config: transformer_2b_config(),
                vae_config: improved_vae_config(),
                scheduler_config: scheduler_config(false),
                base_pass: LtxPassConfig {
                    num_inference_steps: 7,
                    custom_sigmas: Some(LTX_098_DISTILLED_FIRST_PASS_SIGMAS.to_vec()),
                    skip_initial_inference_steps: 0,
                    skip_final_inference_steps: 0,
                    guidance: LtxGuidanceConfig::constant(
                        1.0,
                        0.0,
                        1.0,
                        LTX_098_2B_DISTILLED_SKIP_BLOCKS,
                    ),
                    tone_map_compression_ratio: 0.0,
                },
                decode_timestep: 0.05,
                decode_noise_scale: 0.025,
                mode: LtxPipelineMode::Multiscale,
                multiscale: Some(LtxMultiscaleConfig {
                    downscale_factor: 0.6666666,
                    first_pass: LtxPassConfig {
                        num_inference_steps: 7,
                        custom_sigmas: Some(LTX_098_DISTILLED_FIRST_PASS_SIGMAS.to_vec()),
                        skip_initial_inference_steps: 0,
                        skip_final_inference_steps: 0,
                        guidance: LtxGuidanceConfig::constant(
                            1.0,
                            0.0,
                            1.0,
                            LTX_098_2B_DISTILLED_SKIP_BLOCKS,
                        ),
                        tone_map_compression_ratio: 0.0,
                    },
                    second_pass: LtxPassConfig {
                        num_inference_steps: 3,
                        custom_sigmas: Some(LTX_098_DISTILLED_SECOND_PASS_SIGMAS.to_vec()),
                        skip_initial_inference_steps: 0,
                        skip_final_inference_steps: 0,
                        guidance: LtxGuidanceConfig::constant(
                            1.0,
                            0.0,
                            1.0,
                            LTX_098_2B_DISTILLED_SKIP_BLOCKS,
                        ),
                        tone_map_compression_ratio: 0.0,
                    },
                }),
            })
        } else if model_name.contains("ltx-video-0.9.8-13b-distilled") {
            Ok(Self {
                transformer_config: transformer_13b_config(),
                vae_config: improved_vae_config(),
                scheduler_config: scheduler_config(false),
                base_pass: LtxPassConfig {
                    num_inference_steps: 7,
                    custom_sigmas: Some(LTX_098_DISTILLED_FIRST_PASS_SIGMAS.to_vec()),
                    skip_initial_inference_steps: 0,
                    skip_final_inference_steps: 0,
                    guidance: LtxGuidanceConfig::constant(
                        1.0,
                        0.0,
                        1.0,
                        LTX_098_13B_DISTILLED_SKIP_BLOCKS,
                    ),
                    tone_map_compression_ratio: 0.0,
                },
                decode_timestep: 0.05,
                decode_noise_scale: 0.025,
                mode: LtxPipelineMode::Multiscale,
                multiscale: Some(LtxMultiscaleConfig {
                    downscale_factor: 0.6666666,
                    first_pass: LtxPassConfig {
                        num_inference_steps: 7,
                        custom_sigmas: Some(LTX_098_DISTILLED_FIRST_PASS_SIGMAS.to_vec()),
                        skip_initial_inference_steps: 0,
                        skip_final_inference_steps: 0,
                        guidance: LtxGuidanceConfig::constant(
                            1.0,
                            0.0,
                            1.0,
                            LTX_098_13B_DISTILLED_SKIP_BLOCKS,
                        ),
                        tone_map_compression_ratio: 0.0,
                    },
                    second_pass: LtxPassConfig {
                        num_inference_steps: 3,
                        custom_sigmas: Some(LTX_098_DISTILLED_SECOND_PASS_SIGMAS.to_vec()),
                        skip_initial_inference_steps: 0,
                        skip_final_inference_steps: 0,
                        guidance: LtxGuidanceConfig::constant(
                            1.0,
                            0.0,
                            1.0,
                            LTX_098_13B_DISTILLED_SKIP_BLOCKS,
                        ),
                        tone_map_compression_ratio: 0.6,
                    },
                }),
            })
        } else if model_name.contains("ltx-video-0.9.8-13b-dev") {
            Ok(Self {
                transformer_config: transformer_13b_config(),
                vae_config: improved_vae_config(),
                scheduler_config: scheduler_config(false),
                base_pass: LtxPassConfig {
                    num_inference_steps: 30,
                    custom_sigmas: None,
                    skip_initial_inference_steps: 0,
                    skip_final_inference_steps: 0,
                    guidance: LtxGuidanceConfig::constant(8.0, 4.0, 0.5, &[28]),
                    tone_map_compression_ratio: 0.0,
                },
                decode_timestep: 0.05,
                decode_noise_scale: 0.025,
                mode: LtxPipelineMode::Multiscale,
                multiscale: Some(LtxMultiscaleConfig {
                    downscale_factor: 0.6666666,
                    first_pass: LtxPassConfig {
                        num_inference_steps: 30,
                        custom_sigmas: None,
                        skip_initial_inference_steps: 0,
                        skip_final_inference_steps: 3,
                        guidance: LtxGuidanceConfig {
                            guidance_scale: LTX_098_DEV_FIRST_PASS_GUIDANCE_SCALE.to_vec(),
                            stg_scale: LTX_098_DEV_FIRST_PASS_STG_SCALE.to_vec(),
                            rescaling_scale: LTX_098_DEV_FIRST_PASS_RESCALING_SCALE.to_vec(),
                            guidance_timesteps: Some(
                                LTX_098_DEV_FIRST_PASS_GUIDANCE_TIMESTEPS.to_vec(),
                            ),
                            skip_block_list: LTX_098_13B_DEV_FIRST_PASS_SKIP_BLOCKS
                                .iter()
                                .map(|blocks| blocks.to_vec())
                                .collect(),
                            cfg_star_rescale: true,
                        },
                        tone_map_compression_ratio: 0.0,
                    },
                    second_pass: LtxPassConfig {
                        num_inference_steps: 30,
                        custom_sigmas: None,
                        skip_initial_inference_steps: 17,
                        skip_final_inference_steps: 0,
                        guidance: LtxGuidanceConfig {
                            guidance_scale: vec![1.0],
                            stg_scale: vec![1.0],
                            rescaling_scale: vec![1.0],
                            guidance_timesteps: Some(vec![1.0]),
                            skip_block_list: vec![LTX_098_13B_DEV_SECOND_PASS_SKIP_BLOCKS.to_vec()],
                            cfg_star_rescale: true,
                        },
                        tone_map_compression_ratio: 0.0,
                    },
                }),
            })
        } else {
            bail!("unsupported LTX model preset for {}", model_name);
        }
    }
}

fn transformer_2b_config() -> LtxVideoTransformer3DModelConfig {
    LtxVideoTransformer3DModelConfig {
        num_layers: 28,
        num_attention_heads: 32,
        attention_head_dim: 64,
        cross_attention_dim: 2048,
        caption_channels: 4096,
        ..Default::default()
    }
}

fn transformer_13b_config() -> LtxVideoTransformer3DModelConfig {
    LtxVideoTransformer3DModelConfig {
        num_layers: 48,
        num_attention_heads: 32,
        attention_head_dim: 128,
        cross_attention_dim: 4096,
        caption_channels: 4096,
        ..Default::default()
    }
}

fn improved_vae_config() -> AutoencoderKLLtxVideoConfig {
    AutoencoderKLLtxVideoConfig {
        block_out_channels: vec![128, 256, 512, 1024, 2048],
        decoder_block_out_channels: vec![256, 512, 1024],
        spatiotemporal_scaling: vec![true, true, true, true],
        decoder_spatiotemporal_scaling: vec![true, true, true],
        layers_per_block: vec![4, 6, 6, 2, 2],
        decoder_layers_per_block: vec![5, 5, 5, 5],
        decoder_inject_noise: vec![false, false, false, false],
        decoder_upsample_residual: vec![true, true, true],
        decoder_upsample_factor: vec![2, 2, 2],
        timestep_conditioning: true,
        ..Default::default()
    }
}

fn scheduler_config(stochastic_sampling: bool) -> FlowMatchEulerDiscreteSchedulerConfig {
    FlowMatchEulerDiscreteSchedulerConfig {
        num_train_timesteps: 1000,
        shift: 1.0,
        use_dynamic_shifting: false,
        base_shift: Some(0.5),
        max_shift: Some(1.15),
        base_image_seq_len: Some(256),
        max_image_seq_len: Some(4096),
        invert_sigmas: false,
        shift_terminal: None,
        use_karras_sigmas: false,
        use_exponential_sigmas: false,
        use_beta_sigmas: false,
        time_shift_type: TimeShiftType::Exponential,
        stochastic_sampling,
    }
}

#[derive(Clone, Debug)]
struct LtxResolvedStep {
    guidance_scale: f32,
    stg_scale: f32,
    rescaling_scale: f32,
    skip_blocks: Vec<usize>,
}

fn clamp_skip_blocks(skip_blocks: &[usize], num_layers: usize) -> Vec<usize> {
    skip_blocks
        .iter()
        .copied()
        .filter(|idx| *idx < num_layers)
        .collect()
}

fn resolve_guidance_index(guidance_timesteps: &[f32], sigma: f32) -> usize {
    guidance_timesteps
        .iter()
        .position(|value| *value <= sigma)
        .unwrap_or_else(|| guidance_timesteps.len().saturating_sub(1))
}

fn resolve_step_schedule(
    pass: &LtxPassConfig,
    sigmas: &[f32],
    num_layers: usize,
) -> Vec<LtxResolvedStep> {
    sigmas
        .iter()
        .map(|sigma| {
            let mapped = pass
                .guidance
                .guidance_timesteps
                .as_ref()
                .map(|timesteps| resolve_guidance_index(timesteps, *sigma))
                .unwrap_or(0);
            let value_at = |values: &[f32]| -> f32 {
                if values.len() == 1 {
                    values[0]
                } else {
                    values[mapped.min(values.len() - 1)]
                }
            };
            let skip_blocks = if pass.guidance.skip_block_list.is_empty() {
                Vec::new()
            } else if pass.guidance.skip_block_list.len() == 1 {
                clamp_skip_blocks(&pass.guidance.skip_block_list[0], num_layers)
            } else {
                clamp_skip_blocks(
                    &pass.guidance.skip_block_list
                        [mapped.min(pass.guidance.skip_block_list.len() - 1)],
                    num_layers,
                )
            };
            LtxResolvedStep {
                guidance_scale: value_at(&pass.guidance.guidance_scale),
                stg_scale: value_at(&pass.guidance.stg_scale),
                rescaling_scale: value_at(&pass.guidance.rescaling_scale),
                skip_blocks,
            }
        })
        .collect()
}

fn std_over_dims_except0_keepdim(x: &Tensor) -> Result<Tensor> {
    let rank = x.rank();
    if rank < 2 {
        bail!("std_over_dims_except0_keepdim expects rank >= 2, got {rank}");
    }
    let b = x.dim(0)?;
    let flat = x.flatten_from(1)?;
    let var = flat.var_keepdim(1)?;
    let std = var.sqrt()?;
    let mut shape = Vec::with_capacity(rank);
    shape.push(b);
    shape.extend(std::iter::repeat_n(1usize, rank - 1));
    Ok(std.reshape(shape)?)
}

fn rescale_noise_cfg(
    noise_cfg: &Tensor,
    noise_pred_text: &Tensor,
    guidance_rescale: f32,
) -> Result<Tensor> {
    let std_text = std_over_dims_except0_keepdim(noise_pred_text)?;
    let std_cfg = std_over_dims_except0_keepdim(noise_cfg)?;
    let ratio = std_text.broadcast_div(&std_cfg)?;
    let noise_pred_rescaled = noise_cfg.broadcast_mul(&ratio)?;
    let a = noise_pred_rescaled.affine(guidance_rescale as f64, 0.0)?;
    let b = noise_cfg.affine((1.0 - guidance_rescale) as f64, 0.0)?;
    Ok(a.broadcast_add(&b)?)
}

fn cfg_star_rescale_uncond(noise_pred_uncond: &Tensor, noise_pred_text: &Tensor) -> Result<Tensor> {
    let batch = noise_pred_text.dim(0)?;
    let positive_flat = noise_pred_text.flatten_from(1)?;
    let negative_flat = noise_pred_uncond.flatten_from(1)?;
    let dot = positive_flat
        .broadcast_mul(&negative_flat)?
        .sum_keepdim(1)?;
    let squared = negative_flat.sqr()?.sum_keepdim(1)?.affine(1.0, 1e-8)?;
    let alpha = dot.broadcast_div(&squared)?;
    let alpha = alpha.reshape((batch, 1, 1))?;
    Ok(noise_pred_uncond.broadcast_mul(&alpha.broadcast_as(noise_pred_uncond.shape())?)?)
}

fn create_skip_layer_mask(
    num_layers: usize,
    batch_size: usize,
    layers_to_skip: &[usize],
    device: &Device,
) -> Result<Option<Tensor>> {
    let layers_to_skip = clamp_skip_blocks(layers_to_skip, num_layers);
    if layers_to_skip.is_empty() {
        return Ok(None);
    }

    let mut mask_data = vec![0.0f32; num_layers * batch_size];
    for &layer_idx in &layers_to_skip {
        for batch_idx in 0..batch_size {
            mask_data[layer_idx * batch_size + batch_idx] = 1.0;
        }
    }
    Ok(Some(Tensor::from_vec(
        mask_data,
        (num_layers, batch_size),
        device,
    )?))
}

fn tone_map_latents(latents: &Tensor, compression: f32) -> Result<Tensor> {
    if compression == 0.0 {
        return Ok(latents.clone());
    }
    if !(0.0..=1.0).contains(&compression) {
        bail!("tone map compression must be in [0, 1], got {compression}");
    }
    let scale_factor = compression * 0.75;
    let abs_latents = latents.abs()?;
    let sigmoid_term = abs_latents
        .affine(1.0, -1.0)?
        .affine((4.0 * scale_factor) as f64, 0.0)?;
    let sigmoid_term = candle_nn::ops::sigmoid(&sigmoid_term)?;
    let scales = sigmoid_term.affine((-0.8 * scale_factor) as f64, 1.0)?;
    Ok(latents.broadcast_mul(&scales)?)
}

fn normalize_latents_with_vae(latents: &Tensor, vae: &AutoencoderKLLtxVideo) -> Result<Tensor> {
    let c = latents.dim(1)?;
    let mean = vae
        .latents_mean()
        .reshape((1, c, 1, 1, 1))?
        .to_device(latents.device())?
        .to_dtype(latents.dtype())?;
    let std = vae
        .latents_std()
        .reshape((1, c, 1, 1, 1))?
        .to_device(latents.device())?
        .to_dtype(latents.dtype())?;
    Ok(latents.broadcast_sub(&mean)?.broadcast_div(&std)?)
}

fn denormalize_latents_with_vae(latents: &Tensor, vae: &AutoencoderKLLtxVideo) -> Result<Tensor> {
    let c = latents.dim(1)?;
    let mean = vae
        .latents_mean()
        .reshape((1, c, 1, 1, 1))?
        .to_device(latents.device())?
        .to_dtype(latents.dtype())?;
    let std = vae
        .latents_std()
        .reshape((1, c, 1, 1, 1))?
        .to_device(latents.device())?
        .to_dtype(latents.dtype())?;
    Ok(latents.broadcast_mul(&std)?.broadcast_add(&mean)?)
}

fn adain_filter_latents(latents: &Tensor, reference_latents: &Tensor) -> Result<Tensor> {
    let latents_f32 = latents.to_dtype(DType::F32)?;
    let reference_f32 = reference_latents.to_dtype(DType::F32)?;

    let latents_flat = latents_f32.flatten_from(2)?;
    let reference_flat = reference_f32.flatten_from(2)?;

    let lat_mean = latents_flat.mean_keepdim(2)?;
    let lat_std = latents_flat.var_keepdim(2)?.affine(1.0, 1e-6)?.sqrt()?;
    let ref_mean = reference_flat.mean_keepdim(2)?;
    let ref_std = reference_flat.var_keepdim(2)?.affine(1.0, 1e-6)?.sqrt()?;

    let filtered = latents_flat
        .broadcast_sub(&lat_mean)?
        .broadcast_div(&lat_std)?
        .broadcast_mul(&ref_std)?
        .broadcast_add(&ref_mean)?
        .reshape(latents.shape())?;

    Ok(filtered.to_dtype(latents.dtype())?)
}

// ---------------------------------------------------------------------------
// Loaded state
// ---------------------------------------------------------------------------

#[allow(dead_code)]
struct LoadedLtxVideo {
    transformer: Option<LtxVideoTransformer3DModel>,
    vae: Option<AutoencoderKLLtxVideo>,
    device: Device,
    dtype: DType,
}

// ---------------------------------------------------------------------------
// Engine
// ---------------------------------------------------------------------------

#[allow(dead_code)]
pub struct LtxVideoEngine {
    base: EngineBase<LoadedLtxVideo>,
    t5_variant: Option<String>,
    shared_pool: Option<Arc<Mutex<SharedPool>>>,
}

impl LtxVideoEngine {
    pub fn new(
        model_name: String,
        paths: ModelPaths,
        t5_variant: Option<String>,
        load_strategy: LoadStrategy,
        shared_pool: Option<Arc<Mutex<SharedPool>>>,
    ) -> Self {
        Self {
            base: EngineBase::new(model_name, paths, load_strategy),
            t5_variant,
            shared_pool,
        }
    }
}

// ---------------------------------------------------------------------------
// InferenceEngine trait
// ---------------------------------------------------------------------------

impl crate::engine::InferenceEngine for LtxVideoEngine {
    fn generate(&mut self, req: &GenerateRequest) -> Result<GenerateResponse> {
        let start = Instant::now();
        let preset = LtxModelPreset::for_model(&self.base.model_name)?;

        // Video parameters with defaults
        let num_frames = req.frames.unwrap_or(25);
        let fps = req.fps.unwrap_or(24);
        let steps = req.steps;
        let guidance = req.guidance;

        // Validate frame count: must be 8n+1
        if !(num_frames.wrapping_sub(1)).is_multiple_of(8) {
            bail!(
                "frame count must be 8n+1 (9, 17, 25, 33, ...), got {}",
                num_frames
            );
        }

        let seed = req.seed.unwrap_or_else(rand_seed);
        let width = req.width;
        let height = req.height;

        // Validate dimensions are multiples of 32 (VAE spatial compression)
        if !width.is_multiple_of(VAE_SPATIAL_COMPRESSION as u32)
            || !height.is_multiple_of(VAE_SPATIAL_COMPRESSION as u32)
        {
            bail!(
                "LTX Video requires width and height to be multiples of {}, got {}x{}",
                VAE_SPATIAL_COMPRESSION,
                width,
                height
            );
        }

        // Latent dimensions
        let latent_h = height as usize / VAE_SPATIAL_COMPRESSION;
        let latent_w = width as usize / VAE_SPATIAL_COMPRESSION;
        let latent_f = (num_frames as usize - 1) / VAE_TEMPORAL_COMPRESSION + 1;
        // Always use sequential mode for video (high VRAM usage)
        self.generate_sequential(
            req, &preset, seed, num_frames, fps, steps, guidance, width, height, latent_h,
            latent_w, latent_f, start,
        )
    }

    fn model_name(&self) -> &str {
        &self.base.model_name
    }

    fn is_loaded(&self) -> bool {
        self.base.is_loaded()
    }

    fn load(&mut self) -> Result<()> {
        // Video engine always uses sequential mode — components loaded per-generate
        Ok(())
    }

    fn unload(&mut self) {
        self.base.unload();
    }

    fn set_on_progress(&mut self, callback: ProgressCallback) {
        self.base.set_on_progress(callback);
    }

    fn clear_on_progress(&mut self) {
        self.base.clear_on_progress();
    }

    fn model_paths(&self) -> Option<&ModelPaths> {
        Some(&self.base.paths)
    }
}

// ---------------------------------------------------------------------------
// Sequential generation (load-use-drop each component)
// ---------------------------------------------------------------------------

impl LtxVideoEngine {
    fn load_transformer(
        &self,
        preset: &LtxModelPreset,
        device: &Device,
        dtype: DType,
    ) -> Result<LtxVideoTransformer3DModel> {
        let transformer_files: Vec<std::path::PathBuf> =
            if !self.base.paths.transformer_shards.is_empty() {
                self.base.paths.transformer_shards.clone()
            } else {
                vec![self.base.paths.transformer.clone()]
            };

        let is_gguf = transformer_files
            .first()
            .and_then(|p| p.extension())
            .is_some_and(|e| e == "gguf");
        if is_gguf {
            bail!("GGUF quantized LTX Video transformer is not yet supported — use :bf16 variant");
        }

        // SAFETY: mmap'd safetensors files are immutable model weights.
        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&transformer_files, dtype, device)? };
        let vb = if transformer_files.len() == 1
            && is_official_ltx_transformer_checkpoint(&transformer_files[0])
        {
            vb.rename_f(remap_official_ltx_transformer_key)
        } else {
            vb
        };
        Ok(LtxVideoTransformer3DModel::new(
            &preset.transformer_config,
            vb,
        )?)
    }

    fn load_vae(
        &self,
        preset: &LtxModelPreset,
        device: &Device,
        dtype: DType,
    ) -> Result<AutoencoderKLLtxVideo> {
        // SAFETY: mmap'd safetensors file is immutable model data.
        let vb = unsafe {
            VarBuilder::from_mmaped_safetensors(
                std::slice::from_ref(&self.base.paths.vae),
                dtype,
                device,
            )?
        };
        Ok(AutoencoderKLLtxVideo::new(preset.vae_config.clone(), vb)?)
    }

    #[allow(clippy::too_many_arguments)]
    fn denoise_pass(
        &self,
        transformer: &mut LtxVideoTransformer3DModel,
        prompt_embeds: &Tensor,
        attention_mask: &Tensor,
        uncond_embeds: Option<&Tensor>,
        uncond_mask: Option<&Tensor>,
        pass: &LtxPassConfig,
        scheduler_cfg: &FlowMatchEulerDiscreteSchedulerConfig,
        seed: u64,
        width: u32,
        height: u32,
        num_frames: u32,
        fps: u32,
        device: &Device,
        dtype: DType,
        progress: &crate::progress::ProgressReporter,
        stage_name: &str,
        ltx_debug: bool,
        initial_latents: Option<Tensor>,
    ) -> Result<Tensor> {
        let latent_h = height as usize / VAE_SPATIAL_COMPRESSION;
        let latent_w = width as usize / VAE_SPATIAL_COMPRESSION;
        let latent_f = (num_frames as usize - 1) / VAE_TEMPORAL_COMPRESSION + 1;
        let mut scheduler = FlowMatchEulerDiscreteScheduler::new(scheduler_cfg.clone())?;

        scheduler.set_timesteps(
            if pass.custom_sigmas.is_some() {
                None
            } else {
                Some(pass.num_inference_steps as usize)
            },
            device,
            pass.custom_sigmas.as_deref(),
            None,
            None,
        )?;

        let schedule_sigmas = scheduler
            .sigmas()
            .to_device(&Device::Cpu)?
            .to_vec1::<f32>()?;
        let total_steps = schedule_sigmas.len() - 1;
        if pass.skip_initial_inference_steps + pass.skip_final_inference_steps >= total_steps {
            bail!(
                "invalid LTX pass schedule: skip_initial={} + skip_final={} >= total_steps={}",
                pass.skip_initial_inference_steps,
                pass.skip_final_inference_steps,
                total_steps
            );
        }
        let start_step = pass.skip_initial_inference_steps;
        let end_step = total_steps - pass.skip_final_inference_steps;
        let run_sigmas = schedule_sigmas[start_step..end_step].to_vec();
        scheduler.set_begin_index(start_step);

        let step_schedule =
            resolve_step_schedule(pass, &run_sigmas, transformer.config().num_layers);
        let video_coords = build_video_coords(1, latent_f, latent_h, latent_w, fps, device)?;

        let mut latents = match initial_latents {
            Some(latents) => pack_initial_latents_for_second_pass(&latents)?,
            None => {
                let noise = seeded_randn(
                    seed,
                    &[1, LATENT_CHANNELS, latent_f, latent_h, latent_w],
                    device,
                    DType::F32,
                )?;
                pack_latents(&noise, PATCH_SIZE, PATCH_SIZE_T)?
            }
        };

        progress.stage_start(stage_name);
        let denoise_start = Instant::now();

        for (step, sigma) in run_sigmas.iter().copied().enumerate() {
            let step_start = Instant::now();
            let resolved = &step_schedule[step];
            let batch = latents.dim(0)?;
            let timestep_t = Tensor::full(sigma, (batch,), device)?.to_dtype(dtype)?;
            let latents_input = latents.to_dtype(dtype)?;

            let do_cfg = resolved.guidance_scale > 1.0 && uncond_embeds.is_some();
            let do_stg = resolved.stg_scale > 0.0 && !resolved.skip_blocks.is_empty();

            if do_stg {
                transformer.set_skip_block_list(vec![]);
            } else {
                transformer.set_skip_block_list(resolved.skip_blocks.clone());
            }

            let cond_pred = transformer.forward(
                &latents_input,
                prompt_embeds,
                &timestep_t,
                Some(attention_mask),
                latent_f,
                latent_h,
                latent_w,
                None,
                Some(&video_coords),
                None,
            )?;
            let cond_f32 = cond_pred.to_dtype(DType::F32)?;
            let mut combined = cond_f32.clone();

            if do_cfg {
                let uncond_pred = transformer.forward(
                    &latents_input,
                    uncond_embeds.expect("checked above"),
                    &timestep_t,
                    uncond_mask,
                    latent_f,
                    latent_h,
                    latent_w,
                    None,
                    Some(&video_coords),
                    None,
                )?;
                let mut uncond_f32 = uncond_pred.to_dtype(DType::F32)?;
                if pass.guidance.cfg_star_rescale {
                    uncond_f32 = cfg_star_rescale_uncond(&uncond_f32, &cond_f32)?;
                }
                let diff = cond_f32.broadcast_sub(&uncond_f32)?;
                combined =
                    uncond_f32.broadcast_add(&diff.affine(resolved.guidance_scale as f64, 0.0)?)?;
            }

            if do_stg {
                let skip_layer_mask = create_skip_layer_mask(
                    transformer.config().num_layers,
                    batch,
                    &resolved.skip_blocks,
                    device,
                )?;
                let perturbed = transformer.forward(
                    &latents_input,
                    prompt_embeds,
                    &timestep_t,
                    Some(attention_mask),
                    latent_f,
                    latent_h,
                    latent_w,
                    None,
                    Some(&video_coords),
                    skip_layer_mask.as_ref(),
                )?;
                let perturbed_f32 = perturbed.to_dtype(DType::F32)?;
                let diff_stg = cond_f32.broadcast_sub(&perturbed_f32)?;
                combined =
                    combined.broadcast_add(&diff_stg.affine(resolved.stg_scale as f64, 0.0)?)?;
                if pass.guidance.cfg_star_rescale && resolved.rescaling_scale > 0.0 {
                    combined = rescale_noise_cfg(&combined, &cond_f32, resolved.rescaling_scale)?;
                }
            }

            let model_output =
                if transformer.config().out_channels / 2 == transformer.config().in_channels {
                    combined
                        .chunk(2, 2)?
                        .into_iter()
                        .next()
                        .expect("out_channels / 2 == in_channels implies chunk(2) succeeds")
                } else {
                    combined
                };

            if ltx_debug {
                let out_rms = model_output.sqr()?.mean_all()?.to_scalar::<f32>()?.sqrt();
                let lat_rms = latents.sqr()?.mean_all()?.to_scalar::<f32>()?.sqrt();
                if step < 3 || step == run_sigmas.len() - 1 {
                    progress.info(&format!(
                        "Pass {stage_name} step {}: sigma={:.4}, guidance={:.2}, stg={:.2}, lat_rms={:.4}, out_rms={:.4}",
                        step,
                        sigma,
                        resolved.guidance_scale,
                        resolved.stg_scale,
                        lat_rms,
                        out_rms
                    ));
                }
            }

            latents = scheduler
                .step(&model_output, sigma, &latents, None)?
                .prev_sample;

            progress.emit(ProgressEvent::DenoiseStep {
                step: step + 1,
                total: run_sigmas.len(),
                elapsed: step_start.elapsed(),
            });
        }

        progress.stage_done(stage_name, denoise_start.elapsed());
        unpack_latents(
            &latents,
            latent_f,
            latent_h,
            latent_w,
            PATCH_SIZE,
            PATCH_SIZE_T,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn generate_sequential(
        &mut self,
        req: &GenerateRequest,
        preset: &LtxModelPreset,
        seed: u64,
        num_frames: u32,
        fps: u32,
        steps: u32,
        guidance: f64,
        width: u32,
        height: u32,
        _latent_h: usize,
        _latent_w: usize,
        _latent_f: usize,
        start: Instant,
    ) -> Result<GenerateResponse> {
        let progress = &self.base.progress;
        let paths = &self.base.paths;
        let ltx_debug = std::env::var("MOLD_LTX_DEBUG").is_ok_and(|v| v == "1");

        if preset.mode == LtxPipelineMode::Multiscale && paths.spatial_upscaler.is_none() {
            bail!("LTX 0.9.8 requires a spatial upscaler asset in the pulled model files");
        }

        // Select device
        let device = crate::device::create_device(progress)?;
        let dtype = gpu_dtype(&device);

        progress.info(&format!(
            "LTX Video: {}×{} × {} frames, {} steps, seed {}",
            width, height, num_frames, steps, seed
        ));
        if preset.mode == LtxPipelineMode::Multiscale {
            progress.info("Using the full 0.9.8 multiscale refinement path.");
            if steps != preset.base_pass.num_inference_steps {
                progress.info(&format!(
                    "Ignoring --steps={} for multiscale LTX preset {}; using the preset schedule instead.",
                    steps, self.base.model_name
                ));
            }
            if guidance != preset.base_pass.guidance.guidance_scale[0] as f64 {
                progress.info(&format!(
                    "Ignoring --guidance={guidance:.2} for multiscale LTX preset {}; using the preset guidance schedule instead.",
                    self.base.model_name
                ));
            }
        }

        // ---------------------------------------------------------------
        // Step 1: Encode prompt with T5-XXL
        // ---------------------------------------------------------------
        progress.stage_start("Loading T5-XXL encoder");
        let t5_start = Instant::now();

        let t5_encoder_path = paths
            .t5_encoder
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("T5 encoder path not configured"))?;
        let t5_tokenizer_path = paths
            .t5_tokenizer
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("T5 tokenizer path not configured"))?;

        let mut t5 = crate::encoders::t5::T5Encoder::load(
            t5_encoder_path,
            t5_tokenizer_path,
            &device,
            dtype,
            progress,
        )?;
        progress.stage_done("Loading T5-XXL encoder", t5_start.elapsed());

        progress.stage_start("Encoding prompt");
        let encode_start = Instant::now();
        let prompt_embeds = t5.encode(&req.prompt, &device, dtype)?;
        // prompt_embeds: [1, seq_len, 4096] (T5 encoder already adds batch dim)
        progress.stage_done("Encoding prompt", encode_start.elapsed());

        // Build attention mask (all ones — no padding for single prompt)
        let prompt_seq_len = prompt_embeds.dim(1)?;
        let attention_mask =
            Tensor::ones((1, prompt_seq_len), DType::F32, &device)?.to_dtype(dtype)?;

        let needs_uncond = match preset.mode {
            LtxPipelineMode::Base => guidance > 1.0,
            LtxPipelineMode::Multiscale => preset
                .multiscale
                .as_ref()
                .into_iter()
                .flat_map(|cfg| [&cfg.first_pass, &cfg.second_pass])
                .any(|pass| {
                    pass.guidance
                        .guidance_scale
                        .iter()
                        .any(|scale| *scale > 1.0)
                }),
        };

        let (uncond_embeds, uncond_mask) = if needs_uncond {
            progress.stage_start("Encoding negative prompt (CFG)");
            let ue = t5.encode("", &device, dtype)?;
            let ue_seq = ue.dim(1)?;
            let um = Tensor::ones((1, ue_seq), DType::F32, &device)?.to_dtype(dtype)?;
            progress.stage_done("Encoding negative prompt (CFG)", encode_start.elapsed());
            (Some(ue), Some(um))
        } else {
            (None, None)
        };

        // Drop T5 to free VRAM
        drop(t5);
        device.synchronize()?;
        progress.info("T5 encoder dropped, VRAM freed");

        let (mut latents, decode_width, decode_height, tone_map_compression_ratio) = match preset
            .mode
        {
            LtxPipelineMode::Base => {
                progress.stage_start("Loading LTX Video transformer");
                let transformer_start = Instant::now();
                let mut transformer = self.load_transformer(preset, &device, dtype)?;
                progress.stage_done("Loading LTX Video transformer", transformer_start.elapsed());

                let mut pass = preset.base_pass.clone();
                pass.num_inference_steps = steps;
                pass.guidance.guidance_scale = vec![guidance as f32];

                let latents = self.denoise_pass(
                    &mut transformer,
                    &prompt_embeds,
                    &attention_mask,
                    uncond_embeds.as_ref(),
                    uncond_mask.as_ref().map(|m| m as &Tensor),
                    &pass,
                    &preset.scheduler_config,
                    seed,
                    width,
                    height,
                    num_frames,
                    fps,
                    &device,
                    dtype,
                    progress,
                    "Denoising",
                    ltx_debug,
                    None,
                )?;
                drop(transformer);
                device.synchronize()?;
                (latents, width, height, pass.tone_map_compression_ratio)
            }
            LtxPipelineMode::Multiscale => {
                let multiscale = preset.multiscale.as_ref().expect("multiscale preset");
                let first_width = ((width as f32 * multiscale.downscale_factor) as u32)
                    / VAE_SPATIAL_COMPRESSION as u32
                    * VAE_SPATIAL_COMPRESSION as u32;
                let first_height = ((height as f32 * multiscale.downscale_factor) as u32)
                    / VAE_SPATIAL_COMPRESSION as u32
                    * VAE_SPATIAL_COMPRESSION as u32;

                progress.stage_start("Loading LTX Video transformer");
                let transformer_start = Instant::now();
                let mut first_transformer = self.load_transformer(preset, &device, dtype)?;
                progress.stage_done("Loading LTX Video transformer", transformer_start.elapsed());

                let first_pass_latents = self.denoise_pass(
                    &mut first_transformer,
                    &prompt_embeds,
                    &attention_mask,
                    uncond_embeds.as_ref(),
                    uncond_mask.as_ref().map(|m| m as &Tensor),
                    &multiscale.first_pass,
                    &preset.scheduler_config,
                    seed,
                    first_width,
                    first_height,
                    num_frames,
                    fps,
                    &device,
                    dtype,
                    progress,
                    "Denoising First Pass",
                    ltx_debug,
                    None,
                )?;
                drop(first_transformer);
                device.synchronize()?;

                progress.stage_start("Loading spatial upscaler");
                let spatial_start = Instant::now();
                let vae = self.load_vae(preset, &device, dtype)?;
                let upsampler = LatentUpsampler::load(
                    paths.spatial_upscaler.as_ref().expect("checked above"),
                    dtype,
                    &device,
                )?;
                progress.stage_done("Loading spatial upscaler", spatial_start.elapsed());

                progress.stage_start("Refining multiscale pass");
                let refine_start = Instant::now();
                let first_pass_denorm = cast_latents_for_multiscale_upsampler(
                    &denormalize_latents_with_vae(&first_pass_latents, &vae)?,
                    dtype,
                )?;
                let upsampled_latents =
                    normalize_latents_with_vae(&upsampler.forward(&first_pass_denorm)?, &vae)?;
                let upsampled_latents =
                    adain_filter_latents(&upsampled_latents, &first_pass_latents)?;
                progress.stage_done("Refining multiscale pass", refine_start.elapsed());
                drop(upsampler);
                drop(vae);
                device.synchronize()?;

                let second_width = first_width * 2;
                let second_height = first_height * 2;
                progress.stage_start("Loading LTX Video transformer");
                let transformer_start = Instant::now();
                let mut second_transformer = self.load_transformer(preset, &device, dtype)?;
                progress.stage_done("Loading LTX Video transformer", transformer_start.elapsed());

                let latents = self.denoise_pass(
                    &mut second_transformer,
                    &prompt_embeds,
                    &attention_mask,
                    uncond_embeds.as_ref(),
                    uncond_mask.as_ref().map(|m| m as &Tensor),
                    &multiscale.second_pass,
                    &preset.scheduler_config,
                    seed,
                    second_width,
                    second_height,
                    num_frames,
                    fps,
                    &device,
                    dtype,
                    progress,
                    "Denoising Second Pass",
                    ltx_debug,
                    Some(upsampled_latents),
                )?;
                drop(second_transformer);
                device.synchronize()?;
                (
                    latents,
                    second_width,
                    second_height,
                    multiscale.second_pass.tone_map_compression_ratio,
                )
            }
        };

        // ---------------------------------------------------------------
        // Step 5: Load VAE and decode
        // ---------------------------------------------------------------
        progress.stage_start("Loading VAE decoder");
        let vae_start = Instant::now();
        let vae = self.load_vae(preset, &device, dtype)?;
        progress.stage_done("Loading VAE decoder", vae_start.elapsed());

        progress.stage_start("Decoding video frames");
        let decode_start = Instant::now();

        let decode_timestep = if vae.config().timestep_conditioning {
            if preset.decode_noise_scale > 0.0 {
                let noise =
                    seeded_randn(seed ^ 0xdec0de, latents.shape().dims(), &device, DType::F32)?;
                latents = (&latents * (1.0 - preset.decode_noise_scale as f64))?
                    .broadcast_add(&(noise * preset.decode_noise_scale as f64)?)?;
            }
            Some(Tensor::full(preset.decode_timestep, (1,), &device)?.to_dtype(dtype)?)
        } else {
            None
        };

        latents = tone_map_latents(&latents, tone_map_compression_ratio)?;

        // Un-normalize latents immediately before VAE decode.
        latents = denormalize_latents_with_vae(&latents, &vae)?;

        if ltx_debug {
            let l_f32 = latents.to_dtype(DType::F32)?;
            progress.info(&format!(
                "Latents pre-VAE (un-normalized): mean={:.4}, std={:.4}",
                l_f32.mean_all()?.to_scalar::<f32>()?,
                l_f32.flatten_all()?.var(0)?.to_scalar::<f32>()?.sqrt()
            ));
        }

        latents = latents.to_dtype(dtype)?;
        let (_dec_output, video) = vae.decode(&latents, decode_timestep.as_ref(), false, false)?;
        // video: [B, 3, F, H, W] in model dtype
        if ltx_debug {
            let v_f32 = video.to_dtype(DType::F32)?;
            progress.info(&format!(
                "VAE output: shape={:?}, mean={:.4}, min={:.4}, max={:.4}",
                v_f32.shape(),
                v_f32.mean_all()?.to_scalar::<f32>()?,
                v_f32.flatten_all()?.min(0)?.to_scalar::<f32>()?,
                v_f32.flatten_all()?.max(0)?.to_scalar::<f32>()?
            ));
        }

        progress.stage_done("Decoding video frames", decode_start.elapsed());

        // Drop VAE
        drop(vae);
        device.synchronize()?;

        // ---------------------------------------------------------------
        // Step 6: Post-process and encode video
        // ---------------------------------------------------------------
        // Default to APNG for video output (lossless, metadata-rich)
        let output_format = if req.output_format.is_video() {
            req.output_format
        } else {
            OutputFormat::Apng
        };
        let format_name = output_format.extension().to_uppercase();
        progress.stage_start(&format!("Encoding {format_name}"));
        let encode_start = Instant::now();

        // Convert to [0, 255] u8
        let video = video.to_dtype(DType::F32)?;
        let video = ((video.clamp(-1f32, 1f32)? + 1.0)? * 127.5)?.to_dtype(DType::U8)?;
        let video = video.i(0)?; // Remove batch dim: [3, F, H, W]

        // Extract individual frames as RgbImage
        let num_output_frames = video.dim(1)?;
        let mut frames = Vec::with_capacity(num_output_frames);
        for f in 0..num_output_frames {
            let frame = video.i((.., f, .., ..))?.contiguous()?; // [3, H, W]
            let frame = frame.permute((1, 2, 0))?; // [H, W, 3]
            let frame_data: Vec<u8> = frame.flatten_all()?.to_vec1()?;
            let mut rgb = image::RgbImage::from_raw(decode_width, decode_height, frame_data)
                .ok_or_else(|| anyhow::anyhow!("failed to create frame image"))?;
            if decode_width != width || decode_height != height {
                rgb = image::imageops::resize(
                    &rgb,
                    width,
                    height,
                    image::imageops::FilterType::Triangle,
                );
            }
            frames.push(rgb);
        }

        let video_bytes = match output_format {
            OutputFormat::Apng => {
                let metadata = video_enc::VideoMetadata {
                    prompt: req.prompt.clone(),
                    model: self.base.model_name.clone(),
                    seed,
                    steps,
                    guidance: req.guidance,
                    width,
                    height,
                    frames: num_output_frames as u32,
                    fps,
                };
                video_enc::encode_apng(&frames, fps, Some(&metadata))?
            }
            OutputFormat::Gif => video_enc::encode_gif(&frames, fps)?,
            #[cfg(feature = "webp")]
            OutputFormat::Webp => video_enc::encode_webp(&frames, fps)?,
            #[cfg(feature = "mp4")]
            OutputFormat::Mp4 => video_enc::encode_mp4(&frames, fps)?,
            #[cfg(not(feature = "webp"))]
            OutputFormat::Webp => {
                bail!("WebP output requires the 'webp' feature — rebuild with --features webp")
            }
            #[cfg(not(feature = "mp4"))]
            OutputFormat::Mp4 => {
                bail!("MP4 output requires the 'mp4' feature — rebuild with --features mp4")
            }
            _ => bail!("{format_name} is not a supported video output format"),
        };
        let thumbnail_bytes = video_enc::first_frame_png(&frames)?;
        // Generate a GIF preview only when the caller will use it (TUI gallery or --preview).
        // If the primary format is already GIF, reuse the data; otherwise encode on demand.
        let gif_preview = if req.gif_preview {
            if output_format == OutputFormat::Gif {
                video_bytes.clone()
            } else {
                video_enc::encode_gif(&frames, fps)?
            }
        } else {
            Vec::new()
        };

        progress.stage_done(&format!("Encoding {format_name}"), encode_start.elapsed());

        let generation_time_ms = start.elapsed().as_millis() as u64;
        progress.info(&format!(
            "Done: {} frames, {:.1}s total",
            num_output_frames,
            generation_time_ms as f64 / 1000.0
        ));

        Ok(GenerateResponse {
            images: vec![],
            video: Some(VideoData {
                data: video_bytes,
                format: output_format,
                width,
                height,
                frames: num_output_frames as u32,
                fps,
                thumbnail: thumbnail_bytes,
                gif_preview,
            }),
            generation_time_ms,
            model: self.base.model_name.clone(),
            seed_used: seed,
        })
    }
}

// ---------------------------------------------------------------------------
// Latent packing/unpacking (matches LTX Video pipeline)
// ---------------------------------------------------------------------------

/// Pack latents from [B,C,F,H,W] → [B,S,D] where S = F*H*W, D = C*pt*p*p.
fn pack_latents(latents: &Tensor, patch_size: usize, patch_size_t: usize) -> Result<Tensor> {
    let (b, c, f, h, w) = latents.dims5()?;
    if f % patch_size_t != 0 || h % patch_size != 0 || w % patch_size != 0 {
        bail!("latent dims not divisible by patch sizes");
    }
    let f2 = f / patch_size_t;
    let h2 = h / patch_size;
    let w2 = w / patch_size;

    // [B, C, F2, pt, H2, p, W2, p]
    let x = latents.reshape(&[b, c, f2, patch_size_t, h2, patch_size, w2, patch_size])?;
    // permute → [B, F2, H2, W2, C, pt, p, p]
    let x = x.permute([0, 2, 4, 6, 1, 3, 5, 7])?;
    // flatten last 4 → [B, F2, H2, W2, D]
    let x = x.flatten_from(4)?;
    let d = x.dim(4)?;
    let s = f2 * h2 * w2;
    Ok(x.reshape((b, s, d))?)
}

/// Unpack latents from [B,S,D] → [B,C,F,H,W].
fn unpack_latents(
    latents: &Tensor,
    num_frames: usize,
    height: usize,
    width: usize,
    patch_size: usize,
    patch_size_t: usize,
) -> Result<Tensor> {
    let (b, _s, d) = latents.dims3()?;
    let denom = patch_size_t * patch_size * patch_size;
    if d % denom != 0 {
        bail!("D={d} not divisible by patch product {denom}");
    }
    let c = d / denom;

    let x = latents.reshape(&[
        b,
        num_frames,
        height,
        width,
        c,
        patch_size_t,
        patch_size,
        patch_size,
    ])?;
    // [B, C, F2, pt, H2, p, W2, p]
    let x = x.permute([0, 4, 1, 5, 2, 6, 3, 7])?.contiguous()?;
    Ok(x.reshape((
        b,
        c,
        num_frames * patch_size_t,
        height * patch_size,
        width * patch_size,
    ))?)
}

fn pack_initial_latents_for_second_pass(latents: &Tensor) -> Result<Tensor> {
    pack_latents(&latents.to_dtype(DType::F32)?, PATCH_SIZE, PATCH_SIZE_T)
}

fn cast_latents_for_multiscale_upsampler(latents: &Tensor, dtype: DType) -> Result<Tensor> {
    Ok(latents.to_dtype(dtype)?)
}

/// Build video coordinates for 3D RoPE: [B, seq, 3] with (frame, height, width).
fn build_video_coords(
    batch_size: usize,
    latent_f: usize,
    latent_h: usize,
    latent_w: usize,
    fps: u32,
    device: &Device,
) -> Result<Tensor> {
    let grid_f = Tensor::arange(0u32, latent_f as u32, device)?.to_dtype(DType::F32)?;
    let grid_h = Tensor::arange(0u32, latent_h as u32, device)?.to_dtype(DType::F32)?;
    let grid_w = Tensor::arange(0u32, latent_w as u32, device)?.to_dtype(DType::F32)?;

    let f = grid_f
        .reshape((latent_f, 1, 1))?
        .broadcast_as((latent_f, latent_h, latent_w))?;
    let h = grid_h
        .reshape((1, latent_h, 1))?
        .broadcast_as((latent_f, latent_h, latent_w))?;
    let w = grid_w
        .reshape((1, 1, latent_w))?
        .broadcast_as((latent_f, latent_h, latent_w))?;

    let grid = Tensor::stack(&[f, h, w], 0)?; // [3, F, H, W]
    let seq = latent_f * latent_h * latent_w;
    let grid = grid.flatten_from(1)?.transpose(0, 1)?.unsqueeze(0)?; // [1, seq, 3]

    // Apply compression ratios to coordinates
    let vf = grid.i((.., .., 0))?;
    let vh = grid.i((.., .., 1))?;
    let vw = grid.i((.., .., 2))?;

    // Temporal: (L * 8 + 1 - 8).clamp(0) / fps
    let ts_ratio = VAE_TEMPORAL_COMPRESSION as f64;
    let vf = vf
        .affine(ts_ratio, 1.0 - ts_ratio)?
        .clamp(0.0f32, 10000.0f32)?
        .affine(1.0 / fps as f64, 0.0)?;
    // Spatial: L * 32
    let sp_ratio = VAE_SPATIAL_COMPRESSION as f64;
    let vh = vh.affine(sp_ratio, 0.0)?;
    let vw = vw.affine(sp_ratio, 0.0)?;

    let coords = Tensor::stack(&[vf, vh, vw], candle_core::D::Minus1)?;
    if batch_size > 1 {
        Ok(coords.broadcast_as((batch_size, seq, 3))?)
    } else {
        Ok(coords)
    }
}

#[cfg(test)]
mod tests {
    use super::{
        cast_latents_for_multiscale_upsampler, is_official_ltx_transformer_checkpoint,
        pack_initial_latents_for_second_pass, remap_official_ltx_transformer_key, unpack_latents,
        LtxModelPreset, LtxPipelineMode, LATENT_CHANNELS, LTX_098_DISTILLED_SECOND_PASS_SIGMAS,
        PATCH_SIZE, PATCH_SIZE_T,
    };
    use candle_core::{DType, Device, Tensor};
    use std::path::Path;

    #[test]
    fn detects_official_ltx_single_file_checkpoints() {
        assert!(is_official_ltx_transformer_checkpoint(Path::new(
            "ltxv-2b-0.9.6-distilled-04-25.safetensors"
        )));
        assert!(is_official_ltx_transformer_checkpoint(Path::new(
            "ltxv-13b-0.9.8-dev.safetensors"
        )));
        assert!(!is_official_ltx_transformer_checkpoint(Path::new(
            "diffusion_pytorch_model-00001-of-00002.safetensors"
        )));
        assert!(!is_official_ltx_transformer_checkpoint(Path::new(
            "transformer.gguf"
        )));
    }

    #[test]
    fn remaps_official_transformer_keys_to_upstream_checkpoint_names() {
        assert_eq!(
            remap_official_ltx_transformer_key("proj_in.weight"),
            "model.diffusion_model.patchify_proj.weight"
        );
        assert_eq!(
            remap_official_ltx_transformer_key("time_embed.emb.timestep_embedder.linear_1.weight"),
            "model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.weight"
        );
        assert_eq!(
            remap_official_ltx_transformer_key("transformer_blocks.0.attn1.norm_q.weight"),
            "model.diffusion_model.transformer_blocks.0.attn1.q_norm.weight"
        );
        assert_eq!(
            remap_official_ltx_transformer_key("caption_projection.linear_2.bias"),
            "model.diffusion_model.caption_projection.linear_2.bias"
        );
    }

    #[test]
    fn ltx_presets_only_reference_in_range_skip_blocks() {
        for model_name in [
            "ltx-video-0.9.6:bf16",
            "ltx-video-0.9.6-distilled:bf16",
            "ltx-video-0.9.8-2b-distilled:bf16",
            "ltx-video-0.9.8-13b-dev:bf16",
            "ltx-video-0.9.8-13b-distilled:bf16",
        ] {
            let preset = LtxModelPreset::for_model(model_name).expect("preset should exist");
            let mut all_skip_lists = vec![preset.base_pass.guidance.skip_block_list.clone()];
            if let Some(multiscale) = &preset.multiscale {
                all_skip_lists.push(multiscale.first_pass.guidance.skip_block_list.clone());
                all_skip_lists.push(multiscale.second_pass.guidance.skip_block_list.clone());
            }
            for skip_list_group in all_skip_lists {
                for skip_list in skip_list_group {
                    for skip_block in skip_list {
                        assert!(
                            skip_block < preset.transformer_config.num_layers,
                            "{model_name} skip block {skip_block} is out of range for {} layers",
                            preset.transformer_config.num_layers
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn ltx_098_presets_use_multiscale_mode() {
        for model_name in [
            "ltx-video-0.9.8-2b-distilled:bf16",
            "ltx-video-0.9.8-13b-dev:bf16",
            "ltx-video-0.9.8-13b-distilled:bf16",
        ] {
            let preset = LtxModelPreset::for_model(model_name).expect("preset should exist");
            assert_eq!(preset.mode, LtxPipelineMode::Multiscale, "{model_name}");
            assert!(preset.multiscale.is_some(), "{model_name}");
        }
    }

    #[test]
    fn ltx_098_distilled_second_pass_uses_upstream_sigmas() {
        for model_name in [
            "ltx-video-0.9.8-2b-distilled:bf16",
            "ltx-video-0.9.8-13b-distilled:bf16",
        ] {
            let preset = LtxModelPreset::for_model(model_name).expect("preset should exist");
            let multiscale = preset.multiscale.as_ref().expect("multiscale preset");
            assert_eq!(
                multiscale.second_pass.custom_sigmas.as_deref(),
                Some(LTX_098_DISTILLED_SECOND_PASS_SIGMAS),
                "{model_name}"
            );
        }
    }

    #[test]
    fn multiscale_handoff_normalizes_dtypes_for_upsampler_and_second_pass() {
        let device = Device::Cpu;
        let second_pass_latents =
            Tensor::arange(0f32, (LATENT_CHANNELS * 2 * 4 * 6) as f32, &device)
                .expect("tensor")
                .reshape((1, LATENT_CHANNELS, 2, 4, 6))
                .expect("reshape")
                .to_dtype(DType::BF16)
                .expect("bf16");

        let packed = pack_initial_latents_for_second_pass(&second_pass_latents)
            .expect("second-pass repack should succeed");
        assert_eq!(packed.dtype(), DType::F32);

        let unpacked = unpack_latents(&packed, 2, 4, 6, PATCH_SIZE, PATCH_SIZE_T)
            .expect("unpack should round-trip");
        assert_eq!(unpacked.dtype(), DType::F32);
        assert_eq!(
            unpacked.dims5().expect("dims"),
            (1, LATENT_CHANNELS, 2, 4, 6)
        );
        assert_eq!(
            unpacked
                .flatten_all()
                .expect("flatten")
                .to_vec1::<f32>()
                .expect("vec"),
            second_pass_latents
                .to_dtype(DType::F32)
                .expect("f32")
                .flatten_all()
                .expect("flatten")
                .to_vec1::<f32>()
                .expect("vec")
        );

        let upsampler_input =
            cast_latents_for_multiscale_upsampler(&unpacked, DType::BF16).expect("cast");
        assert_eq!(upsampler_input.dtype(), DType::BF16);
        assert_eq!(
            upsampler_input.dims5().expect("dims"),
            (1, LATENT_CHANNELS, 2, 4, 6)
        );
    }
}
