Model-Optimizer/examples/diffusers/quantization/pipeline_manager.py

# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from collections.abc import Iterator
from typing import Any

import torch
from diffusers import DiffusionPipeline, LTXLatentUpsamplePipeline
from models_utils import MODEL_DEFAULTS, MODEL_PIPELINE, MODEL_REGISTRY, ModelType
from quantize_config import ModelConfig

import modelopt.torch.quantization as mtq


class PipelineManager:
    """Manages diffusion pipeline creation and configuration."""

    def __init__(self, config: ModelConfig, logger: logging.Logger):
        """
        Initialize pipeline manager.

        Args:
            config: Model configuration
            logger: Logger instance
        """
        self.config = config
        self.logger = logger
        self.pipe: Any | None = None
        self.pipe_upsample: LTXLatentUpsamplePipeline | None = None  # For LTX-Video upsampling
        self._transformer: torch.nn.Module | None = None

    @staticmethod
    def create_pipeline_from(
        model_type: ModelType,
        torch_dtype: torch.dtype | dict[str, str | torch.dtype] = torch.bfloat16,
        override_model_path: str | None = None,
    ) -> DiffusionPipeline:
        """
        Create and return an appropriate pipeline based on configuration.

        Returns:
            Configured diffusion pipeline

        Raises:
            ValueError: If model type is unsupported
        """
        try:
            pipeline_cls = MODEL_PIPELINE[model_type]
            if pipeline_cls is None:
                raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.")
            model_id = (
                MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path
            )
            pipe = pipeline_cls.from_pretrained(
                model_id,
                torch_dtype=torch_dtype,
                use_safetensors=True,
                **MODEL_DEFAULTS[model_type].get("from_pretrained_extra_args", {}),
            )
            pipe.set_progress_bar_config(disable=True)
            return pipe
        except Exception as e:
            raise e

    def create_pipeline(self) -> Any:
        """
        Create and return an appropriate pipeline based on configuration.

        Returns:
            Configured diffusion pipeline

        Raises:
            ValueError: If model type is unsupported
        """
        self.logger.info(f"Creating pipeline for {self.config.model_type.value}")
        self.logger.info(f"Model path: {self.config.model_path}")
        self.logger.info(f"Data type: {self.config.model_dtype}")

        try:
            if self.config.model_type == ModelType.LTX2:
                from modelopt.torch.quantization.plugins.diffusion import ltx2 as ltx2_plugin

                ltx2_plugin.register_ltx2_quant_linear()
                self.pipe = self._create_ltx2_pipeline()
                self.logger.info("LTX-2 pipeline created successfully")
                return self.pipe

            pipeline_cls = MODEL_PIPELINE[self.config.model_type]
            if pipeline_cls is None:
                raise ValueError(
                    f"Model type {self.config.model_type.value} does not use diffusers pipelines."
                )
            self.pipe = pipeline_cls.from_pretrained(
                self.config.model_path,
                torch_dtype=self.config.model_dtype,
                use_safetensors=True,
                **MODEL_DEFAULTS[self.config.model_type].get("from_pretrained_extra_args", {}),
            )
            if self.config.model_type == ModelType.LTX_VIDEO_DEV:
                # Optionally load the upsampler pipeline for LTX-Video
                if not self.config.ltx_skip_upsampler:
                    self.logger.info("Loading LTX-Video upsampler pipeline...")
                    self.pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained(
                        "Lightricks/ltxv-spatial-upscaler-0.9.7",
                        vae=self.pipe.vae,
                        torch_dtype=self.config.model_dtype,
                    )
                    self.pipe_upsample.set_progress_bar_config(disable=True)
                else:
                    self.logger.info("Skipping upsampler pipeline for faster calibration")
            self.pipe.set_progress_bar_config(disable=True)

            self.logger.info("Pipeline created successfully")
            return self.pipe

        except Exception as e:
            self.logger.error(f"Failed to create pipeline: {e}")
            raise

    def setup_device(self) -> None:
        """Configure pipeline device placement."""
        if not self.pipe:
            raise RuntimeError("Pipeline not created. Call create_pipeline() first.")

        if self.config.model_type == ModelType.LTX2:
            self.logger.info("Skipping device setup for LTX-2 pipeline (handled internally)")
            return

        if self.config.cpu_offloading:
            self.logger.info("Enabling CPU offloading for memory efficiency")
            self.pipe.enable_model_cpu_offload()
            if self.pipe_upsample:
                self.pipe_upsample.enable_model_cpu_offload()
        else:
            self.logger.info("Moving pipeline to CUDA")
            self.pipe.to("cuda")
            if self.pipe_upsample:
                self.logger.info("Moving upsampler pipeline to CUDA")
                self.pipe_upsample.to("cuda")
        # Enable VAE tiling for LTX-Video to save memory
        if self.config.model_type == ModelType.LTX_VIDEO_DEV:
            if hasattr(self.pipe, "vae") and hasattr(self.pipe.vae, "enable_tiling"):
                self.logger.info("Enabling VAE tiling for LTX-Video")
                self.pipe.vae.enable_tiling()

    def get_backbone(self) -> torch.nn.Module:
        """
        Get the backbone model (transformer or UNet).

        Returns:
            Backbone model module
        """
        if not self.pipe:
            raise RuntimeError("Pipeline not created. Call create_pipeline() first.")

        backbone_pairs = list(self.iter_backbones())
        if len(backbone_pairs) == 1:
            return backbone_pairs[0][1]
        return torch.nn.ModuleList([module for _, module in backbone_pairs])

    def iter_backbones(self) -> Iterator[tuple[str, torch.nn.Module]]:
        """
        Yield backbone modules by name, based on a backbone spec.

        Yields:
            (backbone_name, module) pairs
        """
        if not self.pipe:
            raise RuntimeError("Pipeline not created. Call create_pipeline() first.")

        names = list(self.config.backbone)

        if self.config.model_type == ModelType.LTX2:
            self._ensure_ltx2_transformer_cached()
            name = names[0] if names else "transformer"
            yield name, self._transformer
            return

        if not names:
            raise RuntimeError("No backbone names provided.")

        for name in names:
            module = getattr(self.pipe, name, None)
            if module is None:
                raise RuntimeError(f"Pipeline missing backbone module '{name}'.")
            yield name, module

    def _ensure_ltx2_transformer_cached(self) -> None:
        if not self.pipe:
            raise RuntimeError("Pipeline not created. Call create_pipeline() first.")
        if self._transformer is None:
            transformer = self.pipe.stage_1_model_ledger.transformer()
            self.pipe.stage_1_model_ledger.transformer = lambda: transformer
            self._transformer = transformer

    def _create_ltx2_pipeline(self) -> Any:
        params = dict(self.config.extra_params)
        checkpoint_path = params.pop("checkpoint_path", None)
        distilled_lora_path = params.pop("distilled_lora_path", None)
        distilled_lora_strength = params.pop("distilled_lora_strength", 0.8)
        spatial_upsampler_path = params.pop("spatial_upsampler_path", None)
        gemma_root = params.pop("gemma_root", None)
        fp8_quantization = params.pop("fp8_quantization", None) or params.pop(
            "fp8transformer", False
        )

        if not checkpoint_path:
            raise ValueError("Missing required extra_param: checkpoint_path.")
        if not distilled_lora_path:
            raise ValueError("Missing required extra_param: distilled_lora_path.")
        if not spatial_upsampler_path:
            raise ValueError("Missing required extra_param: spatial_upsampler_path.")
        if not gemma_root:
            raise ValueError("Missing required extra_param: gemma_root.")

        from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
        from ltx_core.quantization import QuantizationPolicy
        from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline

        distilled_lora = [
            LoraPathStrengthAndSDOps(
                str(distilled_lora_path),
                float(distilled_lora_strength),
                LTXV_LORA_COMFY_RENAMING_MAP,
            )
        ]
        pipeline_kwargs = {
            "checkpoint_path": str(checkpoint_path),
            "distilled_lora": distilled_lora,
            "spatial_upsampler_path": str(spatial_upsampler_path),
            "gemma_root": str(gemma_root),
            "loras": [],
            "quantization": QuantizationPolicy.fp8_cast() if fp8_quantization else None,
        }
        pipeline_kwargs.update(params)
        return TI2VidTwoStagesPipeline(**pipeline_kwargs)

    def print_quant_summary(self):
        backbone_pairs = list(self.iter_backbones())
        for name, backbone in backbone_pairs:
            self.logger.info(f"{name} quantization info:")
            mtq.print_quant_summary(backbone)