Model-Optimizer/examples/diffusers/fastgen/preprocess/processors/qwen_image.py

# SPDX-FileCopyrightText: Copyright (c) 2025 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.

"""
Qwen-Image model processor for preprocessing.

Handles Qwen/Qwen-Image T2I models with:
- VAE for image encoding
- Qwen2 text encoder for text conditioning
"""

import logging
from typing import Any

import torch

from .base import BaseModelProcessor
from .registry import ProcessorRegistry

logger = logging.getLogger(__name__)


@ProcessorRegistry.register("qwen_image")
class QwenImageProcessor(BaseModelProcessor):
    """
    Processor for Qwen-Image T2I models.

    Qwen-Image uses a VAE for image encoding and a Qwen2 text encoder
    for text conditioning.
    """

    @property
    def model_type(self) -> str:
        return "qwen_image"

    @property
    def default_model_name(self) -> str:
        return "Qwen/Qwen-Image"

    def load_models(self, model_name: str, device: str) -> dict[str, Any]:
        """
        Load Qwen-Image models.

        Args:
            model_name: HuggingFace model path (e.g., 'Qwen/Qwen-Image')
            device: Device to load models on

        Returns:
            Dict containing:
                - vae: AutoencoderKL
                - tokenizer: Qwen2 tokenizer
                - text_encoder: Qwen2 text encoder
        """
        from diffusers import QwenImagePipeline

        logger.info("[Qwen-Image] Loading models from %s...", model_name)

        # Load pipeline without transformer (not needed for preprocessing)
        pipeline = QwenImagePipeline.from_pretrained(
            model_name,
            transformer=None,
            torch_dtype=torch.bfloat16,
        )

        models = {}

        logger.info("  Configuring VAE...")
        models["vae"] = pipeline.vae.to(device=device, dtype=torch.bfloat16)
        models["vae"].eval()

        logger.info("  Configuring Qwen2 text encoder...")
        pipeline.text_encoder.to(device)
        pipeline.text_encoder.eval()

        # Keep pipeline for encode_prompt — it owns the tokenizer, text_encoder,
        # chat template, and system-token dropping logic.
        models["pipeline"] = pipeline

        torch.cuda.empty_cache()

        logger.info("[Qwen-Image] Models loaded successfully!")
        return models

    def encode_image(
        self,
        image_tensor: torch.Tensor,
        models: dict[str, Any],
        device: str,
    ) -> torch.Tensor:
        """
        Encode image to latent space using VAE.

        Args:
            image_tensor: Image tensor (1, 3, H, W), normalized to [-1, 1]
            models: Dict containing 'vae'
            device: Device to use

        Returns:
            Latent tensor (C, H//8, W//8), FP16
        """
        vae = models["vae"]
        image_tensor = image_tensor.to(device, dtype=torch.bfloat16)

        # Qwen-Image VAE expects 5D input (B, C, T, H, W) — add frame dim for single image
        if image_tensor.ndim == 4:
            image_tensor = image_tensor.unsqueeze(2)

        with torch.no_grad():
            latent = vae.encode(image_tensor).latent_dist.sample()

        # Normalize using per-channel latents_mean / latents_std
        latents_mean = (
            torch.tensor(vae.config.latents_mean)
            .view(1, -1, 1, 1, 1)
            .to(latent.device, latent.dtype)
        )
        latents_std = (
            torch.tensor(vae.config.latents_std)
            .view(1, -1, 1, 1, 1)
            .to(latent.device, latent.dtype)
        )
        latent = (latent - latents_mean) / latents_std

        # Remove frame dim if added, then batch dim → (C, H, W)
        return latent.detach().cpu().to(torch.float16).squeeze(2).squeeze(0)

    def encode_text(
        self,
        prompt: str,
        models: dict[str, Any],
        device: str,
    ) -> dict[str, torch.Tensor]:
        """
        Encode text using the QwenImagePipeline's encode_prompt.

        Delegates to the diffusers pipeline which applies the correct chat
        template, tokenization, system-token dropping, and attention masking.

        Args:
            prompt: Text prompt
            models: Dict containing 'pipeline' (QwenImagePipeline)
            device: Device to use

        Returns:
            Dict containing:
                - prompt_embeds: Qwen2 hidden states [1, seq_len, hidden_dim]
                - prompt_embeds_mask: Attention mask [1, seq_len]
        """
        pipeline = models["pipeline"]

        with torch.no_grad():
            prompt_embeds, prompt_embeds_mask = pipeline.encode_prompt(
                prompt=prompt,
                device=device,
            )

        # Persist the attention mask Qwen-Image's encode_prompt returns (as the docstring
        # promises, and as the negative-prompt path already does). Without it the dataset falls
        # back to an all-ones mask, which is only exact when the cached embeds are trimmed to the
        # real prompt length; keeping the true mask makes the contract explicit and robust.
        encodings = {"prompt_embeds": prompt_embeds.detach().cpu().to(torch.bfloat16)}
        if prompt_embeds_mask is not None:
            encodings["prompt_embeds_mask"] = prompt_embeds_mask.detach().cpu().to(torch.long)
        return encodings

    def verify_latent(
        self,
        latent: torch.Tensor,
        models: dict[str, Any],
        device: str,
    ) -> bool:
        """
        Verify latent can be decoded back to reasonable image.

        Args:
            latent: Encoded latent (C, H, W)
            models: Dict containing 'vae'
            device: Device to use

        Returns:
            True if verification passes
        """
        try:
            vae = models["vae"]

            # (C, H, W) → (B, C, T, H, W) for Qwen-Image VAE
            latent = latent.unsqueeze(0).unsqueeze(2).to(device).float()

            with torch.no_grad():
                # Denormalize: reverse (latent - mean) / std
                latents_mean = (
                    torch.tensor(vae.config.latents_mean)
                    .view(1, -1, 1, 1, 1)
                    .to(device, latent.dtype)
                )
                latents_std = (
                    torch.tensor(vae.config.latents_std)
                    .view(1, -1, 1, 1, 1)
                    .to(device, latent.dtype)
                )
                latent = latent * latents_std + latents_mean
                decoded = vae.decode(latent).sample

            # decoded is 5D (B, C, T, H, W) — take first frame
            decoded = decoded[:, :, 0]
            _, c, h, w = decoded.shape
            if c != 3:
                return False

            return not (torch.isnan(decoded).any() or torch.isinf(decoded).any())

        except Exception as e:
            logger.warning("[Qwen-Image] Verification failed: %s", e)
            return False

    def get_cache_data(
        self,
        latent: torch.Tensor,
        text_encodings: dict[str, torch.Tensor],
        metadata: dict[str, Any],
    ) -> dict[str, Any]:
        """
        Construct cache dictionary for Qwen-Image.

        Args:
            latent: Encoded latent
            text_encodings: Dict from encode_text()
            metadata: Additional metadata

        Returns:
            Dict to save with torch.save()
        """
        cache = {
            # Image latent
            "latent": latent,
            # Text embeddings
            "prompt_embeds": text_encodings["prompt_embeds"],
            # Metadata
            "original_resolution": metadata["original_resolution"],
            "bucket_resolution": metadata["bucket_resolution"],
            "crop_offset": metadata["crop_offset"],
            "prompt": metadata["prompt"],
            "image_path": metadata["image_path"],
            "bucket_id": metadata["bucket_id"],
            "aspect_ratio": metadata["aspect_ratio"],
            # Model info
            "model_type": self.model_type,
        }
        # Carry the positive-prompt attention mask through to the cache when present, so the
        # dataset uses the real mask instead of synthesizing an all-ones one.
        if "prompt_embeds_mask" in text_encodings:
            cache["prompt_embeds_mask"] = text_encodings["prompt_embeds_mask"]
        return cache