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, )