MiniT2I-B-16/pipeline.py

"""Hub custom pipeline: MiniT2ITextToImagePipeline.
Load with native Hugging Face diffusers and trust_remote_code=True.
"""

from __future__ import annotations

from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers
from diffusers.utils import BaseOutput
from diffusers.utils.torch_utils import randn_tensor
# Copyright 2025 The HuggingFace Team. All rights reserved.
#
# 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 inspect
import json
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

os.environ.setdefault("USE_FLAX", "0")
os.environ.setdefault("TRANSFORMERS_NO_FLAX", "1")

import torch
from huggingface_hub import snapshot_download
from PIL import Image
from transformers import AutoTokenizer, T5EncoderModel
from transformers import logging as transformers_logging

transformers_logging.set_verbosity_error()

DEFAULT_NUM_INFERENCE_STEPS = 100
NOISE_INIT_SCALE = 2.0

EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> from pathlib import Path
        >>> import torch
        >>> from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler

        >>> model_dir = Path("./minit2i-diffusers").resolve()
        >>> pipe = DiffusionPipeline.from_pretrained(
        ...     str(model_dir),
        ...     local_files_only=True,
        ...     custom_pipeline=str(model_dir / "pipeline.py"),
        ...     trust_remote_code=True,
        ...     torch_dtype=torch.bfloat16,
        ... )
        >>> pipe.to("cuda")
        >>> pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)

        >>> generator = torch.Generator(device="cuda").manual_seed(42)
        >>> image = pipe(
        ...     "a cinematic portrait of a robot musician",
        ...     num_inference_steps=100,
        ...     guidance_scale=6.0,
        ...     generator=generator,
        ... ).images[0]
        >>> image.save("demo.png")
        ```
"""

MODEL_ALIASES: Dict[str, str] = {
    "b": "minit2i-b-16",
    "b16": "minit2i-b-16",
    "b-16": "minit2i-b-16",
    "base": "minit2i-b-16",
    "minit2i-b16": "minit2i-b-16",
    "minit2i-b-16": "minit2i-b-16",
    "minit2i-b/16": "minit2i-b-16",
    "l": "minit2i-l-16",
    "l16": "minit2i-l-16",
    "l-16": "minit2i-l-16",
    "large": "minit2i-l-16",
    "minit2i-l16": "minit2i-l-16",
    "minit2i-l-16": "minit2i-l-16",
    "minit2i-l/16": "minit2i-l-16",
}

def resolve_model_type(model_type: str) -> str:
    key = model_type.lower().replace("_", "-")
    if key not in MODEL_ALIASES:
        choices = ", ".join(sorted(set(MODEL_ALIASES)))
        raise ValueError(f"Unknown model_type={model_type!r}. Expected one of: {choices}")
    return MODEL_ALIASES[key]

class MiniT2ITextToImagePipeline(DiffusionPipeline):
    r"""
    Text-to-image pipeline for MiniT2I pixel-space flow matching.

    Parameters:
        transformer ([`MiniT2IMMJiTModel`]):
            MiniT2I MM-JiT transformer that predicts flow-matching velocity in pixel space.
        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
            Flow-matching Euler scheduler. Other [`KarrasDiffusionSchedulers`] can be swapped at inference time.
        tokenizer ([`AutoTokenizer`], *optional*):
            Tokenizer for the text encoder.
        text_encoder ([`T5EncoderModel`], *optional*):
            Text encoder used to embed prompts.
    """

    model_cpu_offload_seq = "text_encoder->transformer"
    _optional_components = ["tokenizer", "text_encoder"]

    def __init__(
        self,
        transformer,
        scheduler,
        tokenizer=None,
        text_encoder=None,
        text_encoder_name: str = "google/flan-t5-large",
        model_type: str = "b16",
        repo_id_or_path: Optional[str] = None,
        default_num_inference_steps: int = DEFAULT_NUM_INFERENCE_STEPS,
    ):
        super().__init__()
        if scheduler is None:
            scheduler = self._default_inference_scheduler()
        self.register_modules(
            transformer=transformer,
            scheduler=scheduler,
            tokenizer=tokenizer,
            text_encoder=text_encoder,
        )
        self.register_to_config(
            text_encoder_name=text_encoder_name,
            model_type=model_type,
            repo_id_or_path=repo_id_or_path,
            default_num_inference_steps=int(default_num_inference_steps),
        )
        self._variant_transformers: Dict[str, MiniT2IMMJiTModel] = {}
        self._active_model_type = resolve_model_type(model_type)

    @staticmethod
    def _default_inference_scheduler() -> FlowMatchEulerDiscreteScheduler:
        return FlowMatchEulerDiscreteScheduler(
            num_train_timesteps=1000,
            shift=1.0,
            stochastic_sampling=False,
        )

    @classmethod
    def _load_scheduler_from_dir(
        cls,
        scheduler_dir: Path,
        model_kwargs: Dict[str, Any],
    ) -> Tuple[KarrasDiffusionSchedulers, int]:
        config_path = scheduler_dir / "scheduler_config.json"
        if not config_path.exists():
            return cls._default_inference_scheduler(), DEFAULT_NUM_INFERENCE_STEPS

        config = json.loads(config_path.read_text(encoding="utf-8"))
        class_name = config.get("_class_name", "")
        default_steps = int(config.get("num_inference_steps", DEFAULT_NUM_INFERENCE_STEPS))

        if class_name == "MiniT2IFlowMatchScheduler":
            return cls._default_inference_scheduler(), default_steps

        schedulers_pkg = _hf["schedulers"]
        if hasattr(schedulers_pkg, class_name):
            scheduler_cls = getattr(schedulers_pkg, class_name)
            return scheduler_cls.from_pretrained(str(scheduler_dir), **model_kwargs), default_steps

        return cls._default_inference_scheduler(), default_steps

    @staticmethod
    def _resolve_transformer_path(root: Path, variant_dir: str) -> Path:
        variant_transformer = root / variant_dir / "transformer"
        if variant_transformer.exists():
            return variant_transformer
        root_transformer = root / "transformer"
        if root_transformer.exists():
            return root_transformer
        raise FileNotFoundError(
            f"Could not find transformer weights under {root}. "
            f"Tried {variant_transformer} and {root_transformer}."
        )

    def _get_transformer(
        self,
        model_type: Optional[str],
        repo_id_or_path: Optional[str],
        torch_dtype: Optional[torch.dtype] = None,
        variant: Optional[str] = None,
    ) -> MiniT2IMMJiTModel:
        active_type = resolve_model_type(model_type or self.config.model_type)
        if active_type == self._active_model_type and self.transformer is not None:
            return self.transformer
        if active_type in self._variant_transformers:
            return self._variant_transformers[active_type]

        repo = repo_id_or_path or self.config.repo_id_or_path
        if repo is None:
            raise ValueError("model_type switching requires repo_id_or_path to be set on the pipeline.")

        root = Path(repo)
        if not root.exists():
            root = Path(snapshot_download(repo_id=str(repo)))
        transformer = MiniT2IMMJiTModel.from_pretrained(
            self._resolve_transformer_path(root, active_type),
            torch_dtype=torch_dtype,
            variant=variant,
        )
        self._variant_transformers[active_type] = transformer
        if active_type == resolve_model_type(self.config.model_type):
            self.transformer = transformer
            self._active_model_type = active_type
        return transformer

    @staticmethod
    def prepare_extra_step_kwargs(
        scheduler,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
    ) -> Dict[str, Any]:
        kwargs: Dict[str, Any] = {}
        step_params = set(inspect.signature(scheduler.step).parameters.keys())
        if "generator" in step_params:
            kwargs["generator"] = generator
        return kwargs

    def check_inputs(
        self,
        prompt: Union[str, List[str]],
        guidance_scale: float,
        num_inference_steps: int,
        output_type: str,
    ) -> None:
        if not isinstance(prompt, str) and not (isinstance(prompt, list) and all(isinstance(p, str) for p in prompt)):
            raise TypeError(f"`prompt` must be a string or list of strings, got {type(prompt)}.")
        if guidance_scale < 0:
            raise ValueError(f"`guidance_scale` must be non-negative, got {guidance_scale}.")
        if num_inference_steps <= 0:
            raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}.")
        if output_type not in {"pil", "np", "pt", "latent"}:
            raise ValueError(f"Unsupported `output_type`: {output_type}")

    def prepare_latents(
        self,
        batch_size: int,
        image_size: int,
        in_channels: int,
        device: torch.device,
        dtype: torch.dtype,
        generator: Optional[torch.Generator] = None,
        latents: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        shape = (batch_size, in_channels, image_size, image_size)
        if latents is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
            latents = latents * NOISE_INIT_SCALE
        else:
            latents = latents.to(device=device, dtype=dtype)
            if tuple(latents.shape) != shape:
                raise ValueError(f"Invalid `latents` shape: {tuple(latents.shape)}. Expected {shape}.")
        return latents

    def _encode_prompt(
        self,
        prompt: Union[str, List[str]],
        device: torch.device,
        transformer = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if isinstance(prompt, str):
            prompt = [prompt]
        transformer = transformer or self.transformer
        if self.tokenizer is None:
            self.tokenizer = AutoTokenizer.from_pretrained(self.config.text_encoder_name)
        if self.text_encoder is None:
            self.text_encoder = T5EncoderModel.from_pretrained(self.config.text_encoder_name)
        if next(self.text_encoder.parameters()).device != device:
            self.text_encoder.to(device)
        cfg = transformer.mmjit_config
        tokens = self.tokenizer(
            prompt,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=cfg.prompt_length,
        )
        input_ids = tokens.input_ids.to(device)
        attn = tokens.attention_mask.to(device)
        text = self.text_encoder(input_ids=input_ids, attention_mask=attn).last_hidden_state
        return text, attn

    @staticmethod
    def _cfg_velocity(
        transformer,
        x: torch.Tensor,
        t: torch.Tensor,
        text: torch.Tensor,
        mask: torch.Tensor,
        cfg_scale: float,
    ) -> torch.Tensor:
        batch_size = x.shape[0]
        doubled_x = torch.cat([x, x], dim=0)
        doubled_t = torch.cat([t, t], dim=0)
        doubled_text = torch.cat([text, text], dim=0)
        null_mask = torch.zeros_like(mask)
        doubled_mask = torch.cat([mask, null_mask], dim=0)
        velocity = transformer.pred_velocity(doubled_x, doubled_t, doubled_text, doubled_mask)
        cond, uncond = velocity[:batch_size], velocity[batch_size:]
        cfg_interval = transformer.mmjit_config.cfg_interval
        use_cfg = ((t >= cfg_interval[0]) & (t <= cfg_interval[1])).to(velocity.dtype)
        scale = torch.where(
            use_cfg[:, None, None, None] > 0,
            torch.tensor(cfg_scale, device=x.device, dtype=velocity.dtype),
            torch.tensor(1.0, device=x.device, dtype=velocity.dtype),
        )
        return uncond + (cond - uncond) * scale

    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]],
        num_images_per_prompt: int = 1,
        guidance_scale: float = 6.0,
        num_inference_steps: Optional[int] = None,
        generator: Optional[torch.Generator] = None,
        latents: Optional[torch.Tensor] = None,
        output_type: str = "pil",
        return_dict: bool = True,
        progress: bool = True,
        model_type: Optional[str] = None,
        repo_id_or_path: Optional[str] = None,
        variant: Optional[str] = None,
        torch_dtype: Optional[torch.dtype] = None,
    ) -> Union[ImagePipelineOutput, Tuple]:
        r"""
        Generate images from text prompts with MiniT2I.

        Args:
            prompt (`str` or `list[str]`):
                Text prompt or batch of prompts.
            num_images_per_prompt (`int`, defaults to `1`):
                Number of images to generate per prompt.
            guidance_scale (`float`, defaults to `6.0`):
                Classifier-free guidance scale. CFG is active when `guidance_scale != 1.0`.
            num_inference_steps (`int`, *optional*):
                Number of denoising steps. Defaults to the pipeline config value.
            generator (`torch.Generator`, *optional*):
                RNG for reproducibility.
            latents (`torch.Tensor`, *optional*):
                Pre-generated pixel latents with shape `(batch, channels, height, width)`.
            output_type (`str`, defaults to `"pil"`):
                `"pil"`, `"np"`, `"pt"`, or `"latent"`.
            return_dict (`bool`, defaults to `True`):
                Return [`ImagePipelineOutput`] if True.
            progress (`bool`, defaults to `True`):
                Whether to show a progress bar during denoising.
            model_type (`str`, *optional*):
                MiniT2I variant alias such as `"b16"` or `"l16"`.
            repo_id_or_path (`str`, *optional*):
                Hub id or local path used when switching `model_type`.
            variant (`str`, *optional*):
                Weight variant passed to `from_pretrained`.
            torch_dtype (`torch.dtype`, *optional*):
                Optional dtype override when loading a different transformer variant.
        """
        num_inference_steps = int(num_inference_steps or self.config.default_num_inference_steps)
        self.check_inputs(prompt, guidance_scale, num_inference_steps, output_type)

        transformer = self._get_transformer(model_type, repo_id_or_path, torch_dtype=torch_dtype, variant=variant)
        device = self._execution_device
        transformer = transformer.to(device)

        if isinstance(prompt, str):
            prompt_batch = [prompt] * num_images_per_prompt
        else:
            prompt_batch = []
            for entry in prompt:
                prompt_batch.extend([entry] * num_images_per_prompt)

        batch_size = len(prompt_batch)
        mmjit_cfg = transformer.mmjit_config
        model_dtype = next(transformer.parameters()).dtype

        text, attn = self._encode_prompt(prompt_batch, device, transformer=transformer)
        text = text.to(dtype=model_dtype)
        attn = attn.to(dtype=model_dtype)

        if getattr(self.scheduler.config, "stochastic_sampling", False):
            raise ValueError(
                "MiniT2I expects deterministic FlowMatchEulerDiscreteScheduler stepping "
                "(scheduler.config.stochastic_sampling=False)."
            )

        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator)
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        num_train_timesteps = self.scheduler.config.num_train_timesteps

        latents = self.prepare_latents(
            batch_size=batch_size,
            image_size=mmjit_cfg.image_size,
            in_channels=mmjit_cfg.in_channels,
            device=device,
            dtype=model_dtype,
            generator=generator,
            latents=latents,
        )

        timesteps = self.scheduler.timesteps
        if progress:
            timesteps = self.progress_bar(timesteps)

        using_cfg = guidance_scale != 1.0
        for timestep in timesteps:
            flow_time = 1.0 - float(timestep) / num_train_timesteps
            t = torch.full((batch_size,), flow_time, device=device, dtype=model_dtype)
            if using_cfg:
                velocity = self._cfg_velocity(transformer, latents, t, text, attn, guidance_scale)
            else:
                velocity = transformer.pred_velocity(latents, t, text, attn)

            # MiniT2I integrates velocity from noise (t=0) to data (t=1); flip sign for
            # FlowMatchEulerDiscreteScheduler sigma decreasing from 1 to 0.
            latents = self.scheduler.step(-velocity, timestep, latents, **extra_step_kwargs).prev_sample

        if output_type == "latent":
            images = latents
        else:
            images = (latents.clamp(-1, 1) * 127.5 + 128.0).clamp(0, 255).to(torch.uint8)
            if output_type == "pt":
                images = images.float() / 255.0
            else:
                images = images.permute(0, 2, 3, 1).cpu().numpy()
                if output_type == "pil":
                    images = [Image.fromarray(image) for image in images]

        self.maybe_free_model_hooks()
        if not return_dict:
            return (images,)
        return ImagePipelineOutput(images=images)