f_lite/pipeline.py

import logging
import math
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
from diffusers import AutoencoderKL, DiffusionPipeline
from diffusers.utils import BaseOutput
from diffusers.utils.torch_utils import randn_tensor
from PIL import Image
from torch import FloatTensor
from tqdm.auto import tqdm
from transformers import T5EncoderModel, T5TokenizerFast



logger = logging.getLogger(__name__)


@dataclass
class APGConfig:
    """APG (Augmented Parallel Guidance) configuration"""

    enabled: bool = True
    orthogonal_threshold: float = 0.03


@dataclass
class FLitePipelineOutput(BaseOutput):
    """
    Output class for FLitePipeline pipeline.
    Args:
        images (`List[PIL.Image.Image]` or `np.ndarray`)
            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
    """

    images: Union[List[Image.Image], np.ndarray]


class FLitePipeline(DiffusionPipeline):
    r"""
    Pipeline for text-to-image generation using FLite model.
    This model inherits from [`DiffusionPipeline`].
    """

    model_cpu_offload_seq = "text_encoder->dit_model->vae" 

    dit_model: torch.nn.Module
    vae: AutoencoderKL
    text_encoder: T5EncoderModel
    tokenizer: T5TokenizerFast
    _progress_bar_config: Dict[str, Any]

    def __init__(
        self, dit_model: torch.nn.Module, vae: AutoencoderKL, text_encoder: T5EncoderModel, tokenizer: T5TokenizerFast
    ):
        super().__init__()
        # Register all modules for the pipeline
        # Access DiffusionPipeline's register_modules directly to avoid mypy error
        DiffusionPipeline.register_modules(
            self, dit_model=dit_model, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer
        )

        # Move models to channels last for better performance
        # AutoencoderKL inherits from torch.nn.Module which has these methods
        if hasattr(self.vae, "to"):
            self.vae.to(memory_format=torch.channels_last)
        if hasattr(self.vae, "requires_grad_"):
            self.vae.requires_grad_(False)
        if hasattr(self.text_encoder, "requires_grad_"):
            self.text_encoder.requires_grad_(False)

        # Constants
        self.vae_scale_factor = 8
        self.return_index = -8  # T5 hidden state index to use

    def enable_vae_slicing(self):
        """Enable VAE slicing for memory efficiency."""
        if hasattr(self.vae, "enable_slicing"):
            self.vae.enable_slicing()

    def enable_vae_tiling(self):
        """Enable VAE tiling for memory efficiency."""
        if hasattr(self.vae, "enable_tiling"):
            self.vae.enable_tiling()

    def set_progress_bar_config(self, **kwargs):
        """Set progress bar configuration."""
        self._progress_bar_config = kwargs

    def progress_bar(self, iterable=None, **kwargs):
        """Create progress bar for iterations."""
        self._progress_bar_config = getattr(self, "_progress_bar_config", None) or {}
        config = {**self._progress_bar_config, **kwargs}
        return tqdm(iterable, **config)

    def encode_prompt(
        self,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]] = None,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
        max_sequence_length: int = 512,
        return_index: int = -8,
    ) -> Tuple[FloatTensor, FloatTensor]:
        """Encodes the prompt and negative prompt."""
        if isinstance(prompt, str):
            prompt = [prompt]
        device = device or self.text_encoder.device
        # Text encoder forward pass
        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=max_sequence_length,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids.to(device)
        prompt_embeds = self.text_encoder(text_input_ids, return_dict=True, output_hidden_states=True)
        prompt_embeds_tensor = prompt_embeds.hidden_states[return_index]
        if return_index != -1:
            prompt_embeds_tensor = self.text_encoder.encoder.final_layer_norm(prompt_embeds_tensor)
            prompt_embeds_tensor = self.text_encoder.encoder.dropout(prompt_embeds_tensor)

        dtype = dtype or next(self.text_encoder.parameters()).dtype
        prompt_embeds_tensor = prompt_embeds_tensor.to(dtype=dtype, device=device)

        # Handle negative prompts
        if negative_prompt is None:
            negative_embeds = torch.zeros_like(prompt_embeds_tensor)
        else:
            if isinstance(negative_prompt, str):
                negative_prompt = [negative_prompt]
            negative_result = self.encode_prompt(
                prompt=negative_prompt, device=device, dtype=dtype, return_index=return_index
            )
            negative_embeds = negative_result[0]

        # Explicitly cast both tensors to FloatTensor for mypy
        from typing import cast

        prompt_tensor = cast(FloatTensor, prompt_embeds_tensor.to(dtype=dtype))
        negative_tensor = cast(FloatTensor, negative_embeds.to(dtype=dtype))
        return (prompt_tensor, negative_tensor)

    def to(self, torch_device=None, torch_dtype=None, silence_dtype_warnings=False):
        """Move pipeline components to specified device and dtype."""
        if hasattr(self, "vae"):
            self.vae.to(device=torch_device, dtype=torch_dtype)
        if hasattr(self, "text_encoder"):
            self.text_encoder.to(device=torch_device, dtype=torch_dtype)
        if hasattr(self, "dit_model"):
            self.dit_model.to(device=torch_device, dtype=torch_dtype)
        return self

    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]],
        height: Optional[int] = 1024,
        width: Optional[int] = 1024,
        num_inference_steps: int = 30,
        guidance_scale: float = 3.0,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        dtype: Optional[torch.dtype] = None,
        alpha: Optional[float] = None,
        apg_config: Optional[APGConfig] = None,
        **kwargs,
    ):
        """Generate images from text prompt."""
        batch_size = 1  # TODO: Make this method support batch generation
        # Ensure height and width are not None for calculation
        if height is None:
            height = 1024
        if width is None:
            width = 1024

        dtype = dtype or next(self.dit_model.parameters()).dtype
        apg_config = apg_config or APGConfig()

        device = self._execution_device

        # 2. Encode prompts
        prompt_embeds, negative_embeds = self.encode_prompt(
            prompt=prompt, negative_prompt=negative_prompt, device=self.text_encoder.device, dtype=dtype
        )

        # 3. Initialize latents
        latent_height = height // self.vae_scale_factor
        latent_width = width // self.vae_scale_factor

        if isinstance(generator, list):
            if len(generator) != batch_size:
                raise ValueError(f"Got {len(generator)} generators for {batch_size} samples")

        latents = randn_tensor((batch_size, 16, latent_height, latent_width), generator=generator, device=device, dtype=dtype)
        acc_latents = latents.clone()

        # 4. Calculate alpha if not provided
        if alpha is None:
            image_token_size = latent_height * latent_width
            alpha = 2 * math.sqrt(image_token_size / (64 * 64))

        # 6. Sampling loop
        self.dit_model.eval()
        
        # Check if guidance is needed
        do_classifier_free_guidance = guidance_scale >= 1.0
        
        for i in self.progress_bar(range(num_inference_steps, 0, -1)):
            # Calculate timesteps
            t = i / num_inference_steps
            t_next = (i - 1) / num_inference_steps
            # Scale timesteps according to alpha
            t = t * alpha / (1 + (alpha - 1) * t)
            t_next = t_next * alpha / (1 + (alpha - 1) * t_next)
            dt = t - t_next
            
            # Create tensor with proper device
            t_tensor = torch.tensor([t] * batch_size, device=device, dtype=dtype)
            
            if do_classifier_free_guidance:
                # Duplicate latents for both conditional and unconditional inputs
                latents_input = torch.cat([latents] * 2)
                # Concatenate negative and positive prompt embeddings
                context_input = torch.cat([negative_embeds, prompt_embeds])
                # Duplicate timesteps for the batch
                t_input = torch.cat([t_tensor] * 2)
                
                # Get model predictions in a single pass
                model_outputs = self.dit_model(latents_input, context_input, t_input)
                
                # Split outputs back into unconditional and conditional predictions
                uncond_output, cond_output = model_outputs.chunk(2)
                
                if apg_config.enabled:
                    # Augmented Parallel Guidance
                    dy = cond_output
                    dd = cond_output - uncond_output
                    # Find parallel direction
                    parallel_direction = (dy * dd).sum() / (dy * dy).sum() * dy
                    orthogonal_direction = dd - parallel_direction
                    # Scale orthogonal component
                    orthogonal_std = orthogonal_direction.std()
                    orthogonal_scale = min(1, apg_config.orthogonal_threshold / orthogonal_std)
                    orthogonal_direction = orthogonal_direction * orthogonal_scale
                    model_output = dy + (guidance_scale - 1) * orthogonal_direction
                else:
                    # Standard classifier-free guidance
                    model_output = uncond_output + guidance_scale * (cond_output - uncond_output)
            else:
                # If no guidance needed, just run the model normally
                model_output = self.dit_model(latents, prompt_embeds, t_tensor)
            
            # Update latents
            acc_latents = acc_latents + dt * model_output.to(device)
            latents = acc_latents.clone()

        # 7. Decode latents
        # These checks handle the case where mypy doesn't recognize these attributes
        scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215) if hasattr(self.vae, "config") else 0.18215
        shift_factor = getattr(self.vae.config, "shift_factor", 0) if hasattr(self.vae, "config") else 0

        latents = latents / scaling_factor + shift_factor

        vae_dtype = self.vae.dtype if hasattr(self.vae, "dtype") else dtype
        decoded_images = self.vae.decode(latents.to(vae_dtype)).sample if hasattr(self.vae, "decode") else latents

        # Offload all models
        try:
            self.maybe_free_model_hooks()
        except AttributeError as e:
            if "OptimizedModule" in str(e):
                import warnings
                warnings.warn(
                    "Encountered 'OptimizedModule' error when offloading models. "
                    "This issue might be fixed in the future by: "
                    "https://github.com/huggingface/diffusers/pull/10730"
                )
            else:
                raise
            
        # 8. Post-process images
        images = (decoded_images / 2 + 0.5).clamp(0, 1)
        # Convert to PIL Images
        images = (images * 255).round().clamp(0, 255).to(torch.uint8).cpu()
        pil_images = [Image.fromarray(img.permute(1, 2, 0).numpy()) for img in images]

        return FLitePipelineOutput(
            images=pil_images,
        )