PixelFlow-T2I/pipeline.py

# Copyright 2026 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.

from __future__ import annotations

import inspect

import importlib
import json
import math
import sys
from pathlib import Path
from typing import List, Optional, Tuple, Union, Any

import torch
import torch.nn.functional as F
from einops import rearrange

from diffusers.image_processor import VaeImageProcessor
from diffusers.models.embeddings import get_2d_rotary_pos_embed
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.utils.torch_utils import randn_tensor


DEFAULT_NATIVE_RESOLUTION = 1024

EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> from pathlib import Path
        >>> import sys
        >>> import torch

        >>> model_dir = Path("./PixelFlow-T2I").resolve()
        >>> sys.path.insert(0, str(model_dir))
        >>> from pipeline import PixelFlowT2IPipeline

        >>> pipe = PixelFlowT2IPipeline.from_pretrained(
        ...     str(model_dir),
        ...     local_files_only=True,
        ...     torch_dtype=torch.bfloat16,
        ... )
        >>> pipe.to("cuda")
        ```
"""



class PixelFlowT2IPipeline(DiffusionPipeline):
    r"""
    Pipeline for text-to-image PixelFlow pixel-space cascade generation.

    Parameters:
        transformer ([`PixelFlowTransformer2DModel`]):
            Text-conditioned PixelFlow transformer operating in pixel space.
        scheduler ([`PixelFlowScheduler`]):
            Multi-stage flow scheduler used by PixelFlow.
        text_encoder ([`T5EncoderModel`], *optional*):
            Text encoder used to embed prompts.
        tokenizer ([`T5Tokenizer`], *optional*):
            Tokenizer paired with the text encoder.
    """

    @staticmethod
    def prepare_extra_step_kwargs(
        scheduler,
        generator=None,
        eta: float | None = None,
    ):
        kwargs = {}
        step_params = set(inspect.signature(scheduler.step).parameters.keys())
        if "generator" in step_params:
            kwargs["generator"] = generator
        if eta is not None and "eta" in step_params:
            kwargs["eta"] = eta
        return kwargs


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

    def __init__(
        self,
        transformer: Any,
        scheduler: Any,
        text_encoder=None,
        tokenizer=None,
        max_token_length: int = 512,
    ):
        super().__init__()
        self.register_modules(
            transformer=transformer,
            scheduler=scheduler,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
        )
        self.image_processor = VaeImageProcessor(vae_scale_factor=1, do_normalize=False)
        self.max_token_length = max_token_length

    @staticmethod
    def _prepare_generator(
        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
        device: torch.device,
    ) -> Optional[Union[torch.Generator, List[torch.Generator]]]:
        if generator is None:
            return None
        if isinstance(generator, list):
            return [PixelFlowT2IPipeline._prepare_generator(item, device) for item in generator]

        gen_device = getattr(generator, "device", torch.device("cpu"))
        if gen_device.type == device.type:
            return generator
        new_generator = torch.Generator(device=device)
        seed = int(generator.initial_seed())
        return new_generator.manual_seed(seed)

    def _latent_device(self) -> torch.device:
        transformer_device = getattr(self.transformer, "device", None)
        if transformer_device is not None:
            return transformer_device
        return self._execution_device

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path=None, subfolder=None, **kwargs):
        """Load a self-contained variant folder locally or from the Hub."""
        import importlib
        import sys

        from transformers import T5EncoderModel, T5Tokenizer

        repo_root = Path(__file__).resolve().parent

        if pretrained_model_name_or_path in (None, "", "."):
            variant = repo_root
        elif (
            isinstance(pretrained_model_name_or_path, str)
            and "/" in pretrained_model_name_or_path
            and not Path(pretrained_model_name_or_path).exists()
        ):
            from huggingface_hub import snapshot_download

            hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
            if subfolder:
                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
            cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
            variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
        else:
            variant = Path(pretrained_model_name_or_path)
            if not variant.is_absolute():
                candidate = (Path.cwd() / variant).resolve()
                variant = candidate if candidate.exists() else (repo_root / variant).resolve()
            if subfolder:
                variant = variant / subfolder

        model_kwargs = dict(kwargs)
        scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
        inserted = []

        def _ensure_path(path: str) -> None:
            if path not in sys.path:
                sys.path.insert(0, path)
                inserted.append(path)

        try:
            transformer_dir = variant / "transformer"
            if not (transformer_dir / "transformer_pixelflow.py").exists() or not (transformer_dir / "config.json").exists():
                raise ValueError(f"No loadable transformer found under {variant}")

            _ensure_path(str(transformer_dir))
            transformer_cls = getattr(importlib.import_module("transformer_pixelflow"), "PixelFlowTransformer2DModel")
            transformer = transformer_cls.from_pretrained(str(transformer_dir), **model_kwargs)

            scheduling_py = variant / "scheduling_pixelflow.py"
            scheduler_cfg_dir = variant / "scheduler"
            if not scheduling_py.is_file() or not (scheduler_cfg_dir / "scheduler_config.json").exists():
                raise FileNotFoundError(f"Expected scheduler module at {scheduling_py} and config in {scheduler_cfg_dir}")

            _ensure_path(str(variant.resolve()))
            scheduler_cls = getattr(importlib.import_module("scheduling_pixelflow"), "PixelFlowScheduler")
            try:
                scheduler = scheduler_cls.from_pretrained(str(scheduler_cfg_dir), **scheduler_kwargs)
            except Exception:
                scheduler = scheduler_cls(**scheduler_kwargs)

            text_encoder = None
            tokenizer = None
            text_encoder_dir = variant / "text_encoder"
            tokenizer_dir = variant / "tokenizer"
            if text_encoder_dir.exists() and (text_encoder_dir / "config.json").exists():
                text_encoder = T5EncoderModel.from_pretrained(str(text_encoder_dir), **model_kwargs)
                tokenizer = T5Tokenizer.from_pretrained(str(tokenizer_dir if tokenizer_dir.exists() else text_encoder_dir))

            if text_encoder is None or tokenizer is None:
                text_encoder_name = cls._read_text_encoder_name(variant)
                text_encoder = T5EncoderModel.from_pretrained(text_encoder_name, **model_kwargs)
                tokenizer = T5Tokenizer.from_pretrained(text_encoder_name)

            pipe = cls(transformer=transformer, scheduler=scheduler, text_encoder=text_encoder, tokenizer=tokenizer)
            if hasattr(pipe, "register_to_config"):
                pipe.register_to_config(_name_or_path=str(variant))
            return pipe
        finally:
            for comp_path in inserted:
                if comp_path in sys.path:
                    sys.path.remove(comp_path)

    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        model_kwargs = dict(kwargs)
        transformer_subfolder = model_kwargs.pop("transformer_subfolder", None)
        scheduler_subfolder = model_kwargs.pop("scheduler_subfolder", None)
        text_encoder_subfolder = model_kwargs.pop("text_encoder_subfolder", None)
        tokenizer_subfolder = model_kwargs.pop("tokenizer_subfolder", None)
        scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
        base_path = Path(pretrained_model_name_or_path)

        if transformer_subfolder is None and (base_path / "transformer").exists():
            transformer_subfolder = "transformer"
        if scheduler_subfolder is None and (base_path / "scheduler").exists():
            scheduler_subfolder = "scheduler"
        if text_encoder_subfolder is None and (base_path / "text_encoder").exists():
            text_encoder_subfolder = "text_encoder"
        if tokenizer_subfolder is None and (base_path / "tokenizer").exists():
            tokenizer_subfolder = "tokenizer"

        try:
            return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        except Exception:
            if transformer_subfolder is not None:
                transformer_path = str(base_path / transformer_subfolder)
            else:
                transformer_path = pretrained_model_name_or_path

            transformer = PixelFlowTransformer2DModel.from_pretrained(transformer_path, **model_kwargs)
            try:
                scheduler = PixelFlowScheduler.from_pretrained(
                    pretrained_model_name_or_path,
                    subfolder=scheduler_subfolder,
                    **scheduler_kwargs,
                )
            except Exception:
                scheduler = PixelFlowScheduler(**scheduler_kwargs)

            text_encoder = None
            tokenizer = None
            if text_encoder_subfolder is not None and (base_path / text_encoder_subfolder / "config.json").exists():
                from transformers import T5EncoderModel, T5Tokenizer

                text_encoder = T5EncoderModel.from_pretrained(
                    str(base_path / text_encoder_subfolder),
                    **model_kwargs,
                )
                tokenizer = T5Tokenizer.from_pretrained(str(base_path / tokenizer_subfolder))

            if text_encoder is None and tokenizer is None:
                text_encoder_name = cls._read_text_encoder_name(base_path)
                from transformers import T5EncoderModel, T5Tokenizer

                text_encoder = T5EncoderModel.from_pretrained(text_encoder_name, **model_kwargs)
                tokenizer = T5Tokenizer.from_pretrained(text_encoder_name)

            pipe = cls(
                transformer=transformer,
                scheduler=scheduler,
                text_encoder=text_encoder,
                tokenizer=tokenizer,
            )
            if hasattr(pipe, "register_to_config"):
                pipe.register_to_config(_name_or_path=str(base_path))
            return pipe

    @staticmethod
    def _read_text_encoder_name(variant_path: Path) -> str:
        metadata_path = variant_path / "conversion_metadata.json"
        if metadata_path.exists():
            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
            if metadata.get("text_encoder"):
                return metadata["text_encoder"]
        return "google/flan-t5-xl"

    def check_inputs(
        self,
        prompt: Union[str, List[str]],
        height: int,
        width: int,
        num_inference_steps: Union[int, List[int]],
        output_type: str,
        negative_prompt: Optional[Union[str, List[str]]],
    ) -> None:
        if not isinstance(prompt, str) and not (isinstance(prompt, list) and all(isinstance(p, str) for p in prompt)):
            raise TypeError("`prompt` must be a string or list of strings.")

        if negative_prompt is not None and not isinstance(negative_prompt, str):
            if not (isinstance(negative_prompt, list) and all(isinstance(p, str) for p in negative_prompt)):
                raise TypeError("`negative_prompt` must be a string or list of strings.")

        if output_type not in {"pil", "np", "pt", "latent"}:
            raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.")

        stage_steps = self._normalize_stage_steps(num_inference_steps)
        if any(steps < 1 for steps in stage_steps):
            raise ValueError("Each stage in num_inference_steps must be >= 1.")

        if height <= 0 or width <= 0:
            raise ValueError("height and width must be positive integers.")

    def _normalize_stage_steps(self, num_inference_steps: Union[int, List[int]]) -> List[int]:
        if isinstance(num_inference_steps, int):
            return [num_inference_steps] * self.scheduler.num_stages
        if len(num_inference_steps) != self.scheduler.num_stages:
            raise ValueError(
                f"num_inference_steps must have length {self.scheduler.num_stages} "
                f"(one value per stage), got {len(num_inference_steps)}."
            )
        return list(num_inference_steps)

    def prepare_latents(
        self,
        batch_size: int,
        height: int,
        width: int,
        device: torch.device,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    ) -> Tuple[torch.Tensor, int, int]:
        init_factor = 2 ** (self.scheduler.num_stages - 1)
        coarse_height = height // init_factor
        coarse_width = width // init_factor
        latent_dtype = getattr(self.transformer, "dtype", torch.float32)
        latents = randn_tensor(
            (batch_size, 3, coarse_height, coarse_width),
            generator=generator,
            device=device,
            dtype=latent_dtype,
        )
        return latents, coarse_height, coarse_width

    def _sample_block_noise(
        self,
        batch_size: int,
        channels: int,
        height: int,
        width: int,
        device: torch.device,
        dtype: torch.dtype,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        eps: float = 1e-6,
    ) -> torch.Tensor:
        gamma = self.scheduler.gamma
        cov = torch.eye(4, dtype=torch.float32) * (1 - gamma) + torch.ones(4, 4, dtype=torch.float32) * gamma
        cov = cov + eps * torch.eye(4, dtype=torch.float32)
        chol = torch.linalg.cholesky(cov).to(device=device, dtype=dtype)
        block_number = batch_size * channels * (height // 2) * (width // 2)
        standard = randn_tensor(
            (block_number, 4),
            generator=generator,
            device=device,
            dtype=dtype,
        )
        noise = standard @ chol.T
        return rearrange(
            noise,
            "(b c h w) (p q) -> b c (h p) (w q)",
            b=batch_size,
            c=channels,
            h=height // 2,
            w=width // 2,
            p=2,
            q=2,
        )

    def _upsample_latents_for_stage(
        self,
        latents: torch.Tensor,
        stage_idx: int,
        height: int,
        width: int,
        device: torch.device,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    ) -> torch.Tensor:
        latents = F.interpolate(latents, size=(height, width), mode="nearest")
        original_start_t = self.scheduler.original_start_t[stage_idx]
        gamma = self.scheduler.gamma
        alpha = 1 / (math.sqrt(1 - (1 / gamma)) * (1 - original_start_t) + original_start_t)
        beta = alpha * (1 - original_start_t) / math.sqrt(-gamma)

        noise = self._sample_block_noise(
            *latents.shape,
            device=device,
            dtype=latents.dtype,
            generator=generator,
        )
        return alpha * latents + beta * noise

    def _prepare_rope_pos_embed(self, latents: torch.Tensor, device: torch.device) -> torch.Tensor:
        grid_size = latents.shape[-1] // self.transformer.patch_size
        pos_embed = get_2d_rotary_pos_embed(
            embed_dim=self.transformer.attention_head_dim,
            crops_coords=((0, 0), (grid_size, grid_size)),
            grid_size=(grid_size, grid_size),
            device=device,
            output_type="pt",
        )
        return torch.stack(pos_embed, -1)

    def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"):
        image = (latents / 2 + 0.5).clamp(0, 1)
        if output_type == "latent":
            return latents
        if output_type == "pt":
            return image
        if output_type in {"pil", "np"}:
            return self.image_processor.postprocess(image, output_type=output_type)
        raise ValueError(f"output_type must be one of: 'pil', 'np', 'pt', 'latent'. Got {output_type}.")

    @torch.inference_mode()
    def encode_prompt(
        self,
        prompt: Union[str, List[str]],
        device: torch.device,
        num_images_per_prompt: int = 1,
        do_classifier_free_guidance: bool = True,
        negative_prompt: Union[str, List[str]] = "",
        max_length: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""
        Encode text prompts into hidden states for the PixelFlow transformer.

        Args:
            prompt (`str` or `list[str]`):
                Prompt(s) to encode.
            device (`torch.device`):
                Target device for encoded tensors.
            num_images_per_prompt (`int`, defaults to `1`):
                Number of images to generate per prompt.
            do_classifier_free_guidance (`bool`, defaults to `True`):
                Whether to concatenate unconditional prompt embeddings for CFG.
            negative_prompt (`str` or `list[str]`, defaults to `""`):
                Negative prompt(s) used for classifier-free guidance.
            max_length (`int`, *optional*):
                Maximum token length. Defaults to `self.max_token_length`.
        """
        if self.text_encoder is None or self.tokenizer is None:
            raise ValueError("Text-to-image generation requires `text_encoder` and `tokenizer`.")

        if isinstance(prompt, str):
            prompt = [prompt]
        batch_size = len(prompt)
        max_length = max_length or self.max_token_length

        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=max_length,
            truncation=True,
            add_special_tokens=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids.to(device)
        prompt_attention_mask = text_inputs.attention_mask.to(device)
        prompt_embeds = self.text_encoder(
            text_input_ids,
            attention_mask=prompt_attention_mask,
        )[0]

        dtype = self.text_encoder.dtype
        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
        bs_embed, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
        prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1).repeat(num_images_per_prompt, 1)

        if do_classifier_free_guidance:
            if isinstance(negative_prompt, str):
                uncond_tokens = [negative_prompt] * batch_size
            elif isinstance(negative_prompt, list):
                if len(negative_prompt) != batch_size:
                    raise ValueError(
                        f"Negative prompt list length ({len(negative_prompt)}) must match prompt batch ({batch_size})."
                    )
                uncond_tokens = negative_prompt
            else:
                raise ValueError("Negative prompt must be a string or list of strings.")

            uncond_inputs = self.tokenizer(
                uncond_tokens,
                padding="max_length",
                max_length=prompt_embeds.shape[1],
                truncation=True,
                return_attention_mask=True,
                add_special_tokens=True,
                return_tensors="pt",
            )
            negative_input_ids = uncond_inputs.input_ids.to(device)
            negative_prompt_attention_mask = uncond_inputs.attention_mask.to(device)
            negative_prompt_embeds = self.text_encoder(
                negative_input_ids,
                attention_mask=negative_prompt_attention_mask,
            )[0]

            seq_len_neg = negative_prompt_embeds.shape[1]
            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len_neg, -1)
            negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1).repeat(
                num_images_per_prompt, 1
            )

            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)

        return prompt_embeds, prompt_attention_mask

    @torch.inference_mode()
    def __call__(
        self,
        prompt: Union[str, List[str]],
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: Union[int, List[int]] = 10,
        guidance_scale: float = 4.0,
        shift: float = 1.0,
        negative_prompt: Union[str, List[str]] = "",
        num_images_per_prompt: int = 1,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        output_type: str = "pil",
        return_dict: bool = True,
    ) -> Union[ImagePipelineOutput, Tuple]:
        r"""
        Generate text-to-image samples with PixelFlow.

        Args:
            prompt (`str` or `list[str]`):
                Text prompt(s) describing the desired image.
            height (`int`, *optional*):
                Output image height in pixels. Defaults to the transformer's native resolution.
            width (`int`, *optional*):
                Output image width in pixels. Defaults to the transformer's native resolution.
            num_inference_steps (`int` or `list[int]`, defaults to `10`):
                Number of denoising steps per cascade stage.
            guidance_scale (`float`, defaults to `4.0`):
                Classifier-free guidance scale.
            shift (`float`, defaults to `1.0`):
                Noise shift applied by the scheduler when building stage timesteps.
            negative_prompt (`str` or `list[str]`, defaults to `""`):
                Negative prompt(s) for classifier-free guidance.
            num_images_per_prompt (`int`, defaults to `1`):
                Number of images to generate for each prompt.
            generator (`torch.Generator`, *optional*):
                RNG for reproducibility.
            output_type (`str`, defaults to `"pil"`):
                `"pil"`, `"np"`, `"pt"`, or `"latent"`.
            return_dict (`bool`, defaults to `True`):
                Return [`ImagePipelineOutput`] if True.
        """
        if isinstance(prompt, str):
            prompt_list = [prompt]
        else:
            prompt_list = prompt

        default_size = int(getattr(self.transformer.config, "sample_size", DEFAULT_NATIVE_RESOLUTION))
        height = int(height or default_size)
        width = int(width or default_size)
        self.check_inputs(prompt_list, height, width, num_inference_steps, output_type, negative_prompt)

        device = self._execution_device
        latent_device = self._latent_device()
        do_classifier_free_guidance = guidance_scale > 1.0
        stage_steps = self._normalize_stage_steps(num_inference_steps)
        batch_size = len(prompt_list)
        generator = self._prepare_generator(generator, latent_device)

        prompt_embeds, prompt_attention_mask = self.encode_prompt(
            prompt_list,
            device,
            num_images_per_prompt=num_images_per_prompt,
            do_classifier_free_guidance=do_classifier_free_guidance,
            negative_prompt=negative_prompt,
        )

        latents, height, width = self.prepare_latents(
            batch_size * num_images_per_prompt,
            height,
            width,
            latent_device,
            generator,
        )
        latents = latents.to(device=latent_device, dtype=self.transformer.dtype)
        prompt_embeds = prompt_embeds.to(device=latent_device)
        prompt_attention_mask = prompt_attention_mask.to(device=latent_device)
        size_tensor = torch.tensor([latents.shape[-1] // self.transformer.patch_size], dtype=torch.int32, device=latent_device)

        autocast_enabled = latent_device.type == "cuda"
        autocast_dtype = torch.bfloat16 if autocast_enabled else torch.float32

        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator)

        for stage_idx in range(self.scheduler.num_stages):
            self.scheduler.set_timesteps(stage_steps[stage_idx], stage_idx, device=latent_device, shift=shift)
            timesteps = self.scheduler.Timesteps

            if stage_idx > 0:
                height, width = height * 2, width * 2
                latents = self._upsample_latents_for_stage(
                    latents, stage_idx, height, width, latent_device, generator=generator
                )
                latents = latents.to(dtype=self.transformer.dtype)
                size_tensor = torch.tensor([latents.shape[-1] // self.transformer.patch_size], dtype=torch.int32, device=latent_device)

            rope_pos = self._prepare_rope_pos_embed(latents, latent_device)

            for timestep in timesteps:
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = latent_model_input.to(device=latent_device, dtype=self.transformer.dtype)
                timestep_batch = timestep.expand(latent_model_input.shape[0]).to(device=latent_device, dtype=self.transformer.dtype)
                with torch.autocast(latent_device.type, enabled=autocast_enabled, dtype=autocast_dtype):
                    noise_pred = self.transformer(
                        latent_model_input,
                        encoder_hidden_states=prompt_embeds,
                        encoder_attention_mask=prompt_attention_mask,
                        timestep=timestep_batch,
                        latent_size=size_tensor,
                        pos_embed=rope_pos,
                    ).sample

                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                latents = self.scheduler.step(model_output=noise_pred, sample=latents, **extra_step_kwargs).prev_sample

        image = self.decode_latents(latents, output_type=output_type)
        self.maybe_free_model_hooks()

        if not return_dict:
            return (image,)
        return ImagePipelineOutput(images=image)