dy2000/optimized-diffusers-code

1---2title: "Optimized Diffusers Code"3emoji: 🔥4colorFrom: purple5colorTo: gray6sdk: gradio7sdk_version: 5.31.08app_file: app.py9pinned: false10short_description: 'Optimize Diffusers Code on your hardware.'11---12 13Use an LLM to generate reasonable code snippets in a hardware-aware manner for Diffusers. Still experimental.14 15### Motivation16 17Within the Diffusers, we support a bunch of optimization techniques (refer [here](https://huggingface.co/docs/diffusers/main/en/optimization/memory), [here](https://huggingface.co/docs/diffusers/main/en/optimization/cache), and [here](https://huggingface.co/docs/diffusers/main/en/optimization/fp16)). However, it can be daunting for our users to determine when to use what. Hence, this repository tries to take a stab18at using an LLM to generate reasonable code snippets for a given pipeline checkpoint that respects19user hardware configuration.20 21## Getting started22 23Install the requirements from `requirements.txt`.24 25Configure `GOOGLE_API_KEY` in the environment: `export GOOGLE_API_KEY=...`.26 27Then run:28 29```bash30python e2e_example.py 31```32 33By default, the `e2e_example.py` script uses Flux.1-Dev, but this can be configured through the `--ckpt_id` argument.34 35Full usage:36 37```sh38usage: e2e_example.py [-h] [--ckpt_id CKPT_ID] [--gemini_model GEMINI_MODEL] [--variant VARIANT] [--enable_lossy]39 40options:41  -h, --help            show this help message and exit42  --ckpt_id CKPT_ID     Can be a repo id from the Hub or a local path where the checkpoint is stored.43  --gemini_model GEMINI_MODEL44                        Gemini model to use. Choose from https://ai.google.dev/gemini-api/docs/models.45  --variant VARIANT     If the `ckpt_id` has variants, supply this flag to estimate compute. Example: 'fp16'.46  --enable_lossy        When enabled, the code will include snippets for enabling quantization.47```48 49## Example outputs50 51<details>52<summary>python e2e_example.py (ran on an H100)</summary>53 54````sh55System RAM: 1999.99 GB56RAM Category: large57 58GPU VRAM: 79.65 GB59VRAM Category: large60current_generate_prompt='\npipeline_loading_memory_GB: 31.424\navailable_system_ram_GB: 1999.9855346679688\navailable_gpu_vram_GB: 79.6474609375\nenable_lossy_outputs: False\nenable_torch_compile: True\n'61Sending request to Gemini...62```python63from diffusers import DiffusionPipeline64import torch65 66# User-provided information:67# pipeline_loading_memory_GB: 31.42468# available_system_ram_GB: 1999.9855346679688 (Large RAM)69# available_gpu_vram_GB: 79.6474609375 (Large VRAM)70# enable_lossy_outputs: False71# enable_torch_compile: True72 73# --- Configuration based on user needs and system capabilities ---74 75# Placeholder for the actual checkpoint ID76# Please replace this with your desired model checkpoint ID.77CKPT_ID = "black-forest-labs/FLUX.1-dev" 78 79# Determine dtype. bfloat16 is generally recommended for performance on compatible GPUs.80# Ensure your GPU supports bfloat16 for optimal performance.81dtype = torch.bfloat1682 83# 1. Pipeline Loading and Device Placement:84# Available VRAM (79.64 GB) is significantly greater than the pipeline's loading memory (31.42 GB).85# Therefore, the entire pipeline can comfortably fit and run on the GPU.86print(f"Loading pipeline '{CKPT_ID}' with {dtype} precision...")87pipe = DiffusionPipeline.from_pretrained(CKPT_ID, torch_dtype=dtype)88 89print("Moving pipeline to CUDA (GPU) as VRAM is sufficient...")90pipe = pipe.to("cuda")91 92# 2. Quantization:93# User specified `enable_lossy_outputs: False`, so no quantization is applied.94print("Quantization is NOT applied as per user's preference for lossless outputs.")95 96# 3. Torch Compile:97# User specified `enable_torch_compile: True`.98# Since no offloading was applied (the entire model is on GPU), we can use `fullgraph=True`99# for potentially greater performance benefits.100print("Applying torch.compile() to the transformer for accelerated inference...")101# The transformer is typically the most compute-intensive part of the diffusion pipeline.102# Compiling it can lead to significant speedups.103pipe.transformer.compile(fullgraph=True)104 105# --- Inference ---106print("Starting inference...")107prompt = "photo of a dog sitting beside a river, high quality, 4k"108image = pipe(prompt).images[0]109 110print("Inference completed. Displaying image.")111# Save or display the image112image.save("generated_image.png")113print("Image saved as generated_image.png")114 115# You can also display the image directly if running in an environment that supports it116# image.show()117```118````119<br>120</details>121<br>122<details>123<summary>python e2e_example.py --enable_lossy</summary>124 125````sh126System RAM: 1999.99 GB127RAM Category: large128 129GPU VRAM: 79.65 GB130VRAM Category: large131current_generate_prompt='\npipeline_loading_memory_GB: 31.424\navailable_system_ram_GB: 1999.9855346679688\navailable_gpu_vram_GB: 79.6474609375\nenable_lossy_outputs: True\nenable_torch_compile: True\n'132Sending request to Gemini...133```python134import torch135from diffusers import DiffusionPipeline136from diffusers.quantizers import PipelineQuantizationConfig137import os138 139# --- User-provided information and derived constants ---140# Checkpoint ID (assuming a placeholder since it was not provided in the user input)141# Using the example CKPT_ID from the problem description142CKPT_ID = "black-forest-labs/FLUX.1-dev"143 144# Derived from available_gpu_vram_GB (79.64 GB) and pipeline_loading_memory_GB (31.424 GB)145# VRAM is ample to load the entire pipeline146use_cuda_direct_load = True 147 148# Derived from enable_lossy_outputs (True)149enable_quantization = True150 151# Derived from enable_torch_compile (True)152enable_torch_compile = True153 154# --- Inference Code ---155 156print(f"Loading pipeline: {CKPT_ID}")157 158# 1. Quantization Configuration (since enable_lossy_outputs is True)159quant_config = None160if enable_quantization:161    # Default to bitsandbytes 4-bit as per guidance162    print("Enabling bitsandbytes 4-bit quantization for 'transformer' component.")163    quant_config = PipelineQuantizationConfig(164        quant_backend="bitsandbytes_4bit", 165        quant_kwargs={"load_in_4bit": True, "bnb_4bit_compute_dtype": torch.bfloat16, "bnb_4bit_quant_type": "nf4"},166        # For FLUX.1-dev, the main generative component is typically 'transformer'.167        # For other pipelines, you might include 'unet', 'text_encoder', 'text_encoder_2', etc.168        components_to_quantize=["transformer"] 169    )170 171# 2. Load the Diffusion Pipeline172# Use bfloat16 for better performance and modern GPU compatibility173pipe = DiffusionPipeline.from_pretrained(174    CKPT_ID, 175    torch_dtype=torch.bfloat16,176    quantization_config=quant_config if enable_quantization else None177)178 179# 3. Move Pipeline to GPU (since VRAM is ample)180if use_cuda_direct_load:181    print("Moving the entire pipeline to CUDA (GPU).")182    pipe = pipe.to("cuda")183 184# 4. Apply torch.compile() (since enable_torch_compile is True)185if enable_torch_compile:186    print("Applying torch.compile() for speedup.")187    # This setting is beneficial when bitsandbytes is used188    torch._dynamo.config.capture_dynamic_output_shape_ops = True 189    190    # Since no offloading is applied (model fits fully in VRAM), use fullgraph=True191    # The primary component for compilation in FLUX.1-dev is 'transformer'192    print("Compiling pipe.transformer with fullgraph=True.")193    pipe.transformer = torch.compile(pipe.transformer, fullgraph=True)194 195# 5. Perform Inference196print("Starting image generation...")197prompt = "photo of a dog sitting beside a river"198num_inference_steps = 28 # A reasonable number of steps for good quality199 200# Ensure all inputs are on the correct device for inference after compilation201with torch.no_grad():202    image = pipe(prompt, num_inference_steps=num_inference_steps).images[0]203 204print("Image generation complete.")205# Save or display the image206output_path = "generated_image.png"207image.save(output_path)208print(f"Image saved to {output_path}")209 210```211````212 213</details>214<br>215When invoked from an RTX 4090, it outputs:216 217<details>218<summary>Expand</summary>219 220````sh221System RAM: 125.54 GB222RAM Category: large223 224GPU VRAM: 23.99 GB225VRAM Category: medium226current_generate_prompt='\npipeline_loading_memory_GB: 31.424\navailable_system_ram_GB: 125.54026794433594\navailable_gpu_vram_GB: 23.98828125\nenable_lossy_outputs: False\nenable_torch_compile: True\n'227Sending request to Gemini...228```python229import torch230from diffusers import DiffusionPipeline231import os # For creating offload directories if needed, though not directly used in this solution232 233# --- User-provided information (interpreted) ---234# Checkpoint ID will be a placeholder as it's not provided by the user directly in the input.235# pipeline_loading_memory_GB: 31.424 GB236# available_system_ram_GB: 125.54 GB (Categorized as "large": > 40GB)237# available_gpu_vram_GB: 23.98 GB (Categorized as "medium": > 8GB <= 24GB)238# enable_lossy_outputs: False (User prefers no quantization)239# enable_torch_compile: True (User wants to enable torch.compile)240 241# --- Configuration ---242# Placeholder for the actual checkpoint ID. Replace with the desired model ID.243CKPT_ID = "black-forest-labs/FLUX.1-dev" # Example from Diffusers library.244PROMPT = "photo of a dog sitting beside a river"245 246print(f"--- Optimizing inference for CKPT_ID: {CKPT_ID} ---")247print(f"Pipeline loading memory: {31.424} GB")248print(f"Available System RAM: {125.54} GB (Large)")249print(f"Available GPU VRAM: {23.98} GB (Medium)")250print(f"Lossy outputs (quantization): {'Disabled' if not False else 'Enabled'}")251print(f"Torch.compile: {'Enabled' if True else 'Disabled'}")252print("-" * 50)253 254# --- 1. Load the Diffusion Pipeline ---255# Use bfloat16 for a good balance of memory and performance.256print(f"Loading pipeline '{CKPT_ID}' with torch_dtype=torch.bfloat16...")257pipe = DiffusionPipeline.from_pretrained(CKPT_ID, torch_dtype=torch.bfloat16)258print("Pipeline loaded.")259 260# --- 2. Apply Memory Optimizations ---261# Analysis:262# - Pipeline memory (31.424 GB) exceeds available GPU VRAM (23.98 GB).263# - System RAM (125.54 GB) is large.264# Strategy: Use `enable_model_cpu_offload()`. This moves model components to CPU when not265# in use, swapping them to GPU on demand. This is ideal when VRAM is insufficient but system266# RAM is abundant.267 268print("Applying memory optimization: `pipe.enable_model_cpu_offload()`...")269pipe.enable_model_cpu_offload()270print("Model CPU offloading enabled. Components will dynamically move between CPU and GPU.")271 272# --- 3. Apply Speed Optimizations (torch.compile) ---273# Analysis:274# - `enable_torch_compile` is True.275# - Model offloading (`enable_model_cpu_offload`) is applied.276# Strategy: Enable torch.compile with `recompile_limit` as offloading is used.277# Do not use `fullgraph=True` when offloading is active.278 279print("Applying speed optimization: `torch.compile()`...")280torch._dynamo.config.recompile_limit = 1000 # Recommended when offloading is applied.281# torch._dynamo.config.capture_dynamic_output_shape_ops = True # Only for bitsandbytes, not applicable here.282 283# Compile the main computational component (e.g., transformer or unet).284# FLUX models primarily use a transformer. For other models, it might be `pipe.unet`.285if hasattr(pipe, "transformer"):286    print("Compiling `pipe.transformer`...")287    pipe.transformer.compile()288elif hasattr(pipe, "unet"):289    print("Compiling `pipe.unet`...")290    pipe.unet.compile()291else:292    print("Warning: Neither `pipe.transformer` nor `pipe.unet` found for compilation. Skipping `torch.compile` for core component.")293 294print("Speed optimizations applied.")295 296# --- 4. Perform Inference ---297print(f"Starting image generation for prompt: '{PROMPT}'")298# The pipeline handles component swapping automatically due to `enable_model_cpu_offload()`.299image = pipe(PROMPT).images[0]300print("Image generation complete.")301 302# --- 5. Save or Display the Result ---303output_path = "generated_image.png"304image.save(output_path)305print(f"Generated image saved to '{output_path}'")306 307print("\n--- Inference process finished successfully ---")308```309````310</details>311 312### More outputs313 314<details>315<summary>"Wan-AI/Wan2.1-T2V-14B-Diffusers" with lossy outputs enabled</summary>316 317````sh318System RAM: 125.54 GB319RAM Category: large320 321GPU VRAM: 23.99 GB322VRAM Category: medium323("current_generate_prompt='\\nckpt_id: "324 'Wan-AI/Wan2.1-T2V-14B-Diffusers\\npipeline_loading_memory_GB: '325 '37.432\\navailable_system_ram_GB: '326 '125.54026794433594\\navailable_gpu_vram_GB: '327 '23.98828125\\nenable_lossy_outputs: True\\nis_fp8_supported: '328 "True\\nenable_torch_compile: True\\n'")329Sending request to Gemini...330```python331from diffusers import DiffusionPipeline332from diffusers.quantizers import PipelineQuantizationConfig333import torch334 335ckpt_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"336 337quant_config = PipelineQuantizationConfig(338    quant_backend="torchao",339    quant_kwargs={"quant_type": "float8dq_e4m3_row"},340    components_to_quantize=["transformer"]341)342pipe = DiffusionPipeline.from_pretrained(ckpt_id, quantization_config=quant_config, torch_dtype=torch.bfloat16)343 344# Apply model CPU offload due to VRAM constraints345pipe.enable_model_cpu_offload()346 347# torch.compile() configuration348torch._dynamo.config.recompile_limit = 1000349pipe.transformer.compile()350# pipe.vae.decode = torch.compile(pipe.vae.decode) # Uncomment if you want to compile VAE decode as well351 352prompt = "photo of a dog sitting beside a river"353 354# Modify the pipe call arguments as needed.355image = pipe(prompt).images[0]356 357# You can save the image or perform further operations here358# image.save("generated_image.png")359```360````361</details>362<small>Ran on an RTX 4090</small>