Merge branch 'main' into save-load-optional-components-tests

.github/workflows/nightly_tests.yml

@@ -418,8 +418,6 @@ jobs:
             test_location: "gguf"
           - backend: "torchao"
             test_location: "torchao"
-          - backend: "optimum_quanto"
-            test_location: "quanto"
     runs-on:
       group: aws-g6e-xlarge-plus
     container:

.github/workflows/pr_tests.yml

@@ -3,6 +3,7 @@ name: Fast tests for PRs
 on:
   pull_request:
     branches: [main]
+    types: [synchronize]
     paths:
       - "src/diffusers/**.py"
       - "benchmarks/**.py"

docs/source/en/_toctree.yml

@@ -173,8 +173,6 @@
     title: gguf
   - local: quantization/torchao
     title: torchao
-  - local: quantization/quanto 
-    title: quanto
   title: Quantization Methods
 - sections:
   - local: optimization/fp16

docs/source/en/api/quantization.md

@@ -31,11 +31,6 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui
 ## GGUFQuantizationConfig
 
 [[autodoc]] GGUFQuantizationConfig
-
-## QuantoConfig
-
-[[autodoc]] QuantoConfig
-
 ## TorchAoConfig
 
 [[autodoc]] TorchAoConfig

docs/source/en/quantization/overview.md

@@ -36,6 +36,5 @@ Diffusers currently supports the following quantization methods.
 - [BitsandBytes](./bitsandbytes)
 - [TorchAO](./torchao)
 - [GGUF](./gguf)
-- [Quanto](./quanto.md)
 
 [This resource](https://huggingface.co/docs/transformers/main/en/quantization/overview#when-to-use-what) provides a good overview of the pros and cons of different quantization techniques.

docs/source/en/quantization/quanto.md

@@ -1,148 +0,0 @@
-<!--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.
-
--->
-
-# Quanto
-
-[Quanto](https://github.com/huggingface/optimum-quanto) is a PyTorch quantization backend for [Optimum](https://huggingface.co/docs/optimum/en/index). It has been designed with versatility and simplicity in mind:
-
-- All features are available in eager mode (works with non-traceable models)
-- Supports quantization aware training
-- Quantized models are compatible with `torch.compile`
-- Quantized models are Device agnostic (e.g CUDA,XPU,MPS,CPU)
-
-In order to use the Quanto backend, you will first need to install `optimum-quanto>=0.2.6` and `accelerate`
-
-```shell
-pip install optimum-quanto accelerate
-```
-
-Now you can quantize a model by passing the `QuantoConfig` object to the `from_pretrained()` method. Although the Quanto library does allow quantizing `nn.Conv2d` and `nn.LayerNorm` modules, currently, Diffusers only supports quantizing the weights in the `nn.Linear` layers of a model. The following snippet demonstrates how to apply `float8` quantization with Quanto.   
-
-```python
-import torch
-from diffusers import FluxTransformer2DModel, QuantoConfig
-
-model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights_dtype="float8")
-transformer = FluxTransformer2DModel.from_pretrained(
-      model_id,
-      subfolder="transformer",
-      quantization_config=quantization_config,
-      torch_dtype=torch.bfloat16,
-)
-
-pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch_dtype)
-pipe.to("cuda")
-
-prompt = "A cat holding a sign that says hello world"
-image = pipe(
-    prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
-).images[0]
-image.save("output.png")
-```
-
-## Skipping Quantization on specific modules
-
-It is possible to skip applying quantization on certain modules using the `modules_to_not_convert` argument in the `QuantoConfig`. Please ensure that the modules passed in to this argument match the keys of the modules in the `state_dict`  
-
-```python
-import torch
-from diffusers import FluxTransformer2DModel, QuantoConfig
-
-model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights_dtype="float8", modules_to_not_convert=["proj_out"])
-transformer = FluxTransformer2DModel.from_pretrained(
-      model_id,
-      subfolder="transformer",
-      quantization_config=quantization_config,
-      torch_dtype=torch.bfloat16,
-)
-```
-
-## Using `from_single_file` with the Quanto Backend
-
-`QuantoConfig` is compatible with `~FromOriginalModelMixin.from_single_file`. 
-
-```python
-import torch
-from diffusers import FluxTransformer2DModel, QuantoConfig
-
-ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
-quantization_config = QuantoConfig(weights_dtype="float8")
-transformer = FluxTransformer2DModel.from_single_file(ckpt_path, quantization_config=quantization_config, torch_dtype=torch.bfloat16)
-```
-
-## Saving Quantized models
-
-Diffusers supports serializing Quanto models using the `~ModelMixin.save_pretrained` method.
-
-The serialization and loading requirements are different for models quantized directly with the Quanto library and models quantized
-with Diffusers using Quanto as the backend. It is currently not possible to load models quantized directly with Quanto into Diffusers using `~ModelMixin.from_pretrained`
-
-```python
-import torch
-from diffusers import FluxTransformer2DModel, QuantoConfig
-
-model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights_dtype="float8")
-transformer = FluxTransformer2DModel.from_pretrained(
-      model_id,
-      subfolder="transformer",
-      quantization_config=quantization_config,
-      torch_dtype=torch.bfloat16,
-)
-# save quantized model to reuse
-transformer.save_pretrained("<your quantized model save path>")
-
-# you can reload your quantized model with
-model = FluxTransformer2DModel.from_pretrained("<your quantized model save path>")
-```
-
-## Using `torch.compile` with Quanto
-
-Currently the Quanto backend supports `torch.compile` for the following quantization types:
-
-- `int8` weights 
-
-```python
-import torch
-from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig
-
-model_id = "black-forest-labs/FLUX.1-dev"
-quantization_config = QuantoConfig(weights_dtype="int8")
-transformer = FluxTransformer2DModel.from_pretrained(
-    model_id,
-    subfolder="transformer",
-    quantization_config=quantization_config,
-    torch_dtype=torch.bfloat16,
-)
-transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
-
-pipe = FluxPipeline.from_pretrained(
-    model_id, transformer=transformer, torch_dtype=torch_dtype
-)
-pipe.to("cuda")
-images = pipe("A cat holding a sign that says hello").images[0]
-images.save("flux-quanto-compile.png")
-```
-
-## Supported Quantization Types
-
-### Weights
-
-- float8
-- int8
-- int4
-- int2
-
-

examples/community/README.md

@@ -83,7 +83,6 @@ PIXART-α Controlnet pipeline | Implementation of the controlnet model for pixar
 | [🪆Matryoshka Diffusion Models](https://huggingface.co/papers/2310.15111) | A diffusion process that denoises inputs at multiple resolutions jointly and uses a NestedUNet architecture where features and parameters for small scale inputs are nested within those of the large scales. See [original codebase](https://github.com/apple/ml-mdm). | [🪆Matryoshka Diffusion Models](#matryoshka-diffusion-models) | [![Hugging Face Space](https://img.shields.io/badge/🤗%20Hugging%20Face-Space-yellow)](https://huggingface.co/spaces/pcuenq/mdm) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/gist/tolgacangoz/1f54875fc7aeaabcf284ebde64820966/matryoshka_hf.ipynb) | [M. Tolga Cangöz](https://github.com/tolgacangoz) |
 | Stable Diffusion XL Attentive Eraser Pipeline |[[AAAI2025 Oral] Attentive Eraser](https://github.com/Anonym0u3/AttentiveEraser) is a novel tuning-free method that enhances object removal capabilities in pre-trained diffusion models.|[Stable Diffusion XL Attentive Eraser Pipeline](#stable-diffusion-xl-attentive-eraser-pipeline)|-|[Wenhao Sun](https://github.com/Anonym0u3) and [Benlei Cui](https://github.com/Benny079)|
 | Perturbed-Attention Guidance |StableDiffusionPAGPipeline is a modification of StableDiffusionPipeline to support Perturbed-Attention Guidance (PAG).|[Perturbed-Attention Guidance](#perturbed-attention-guidance)|[Notebook](https://github.com/huggingface/notebooks/blob/main/diffusers/perturbed_attention_guidance.ipynb)|[Hyoungwon Cho](https://github.com/HyoungwonCho)|
-| CogVideoX DDIM Inversion Pipeline | Implementation of DDIM inversion and guided attention-based editing denoising process on CogVideoX. | [CogVideoX DDIM Inversion Pipeline](#cogvideox-ddim-inversion-pipeline) | - | [LittleNyima](https://github.com/LittleNyima) |
 
 To load a custom pipeline you just need to pass the `custom_pipeline` argument to `DiffusionPipeline`, as one of the files in `diffusers/examples/community`. Feel free to send a PR with your own pipelines, we will merge them quickly.
 
@@ -5223,39 +5222,3 @@ with torch.no_grad():
 
 In the folder examples/pixart there is also a script that can be used to train new models.
 Please check the script `train_controlnet_hf_diffusers.sh` on how to start the training.
-
-# CogVideoX DDIM Inversion Pipeline
-
-This implementation performs DDIM inversion on the video based on CogVideoX and uses guided attention to reconstruct or edit the inversion latents.
-
-## Example Usage
-
-```python
-import torch
-
-from examples.community.cogvideox_ddim_inversion import CogVideoXPipelineForDDIMInversion
-
-
-# Load pretrained pipeline
-pipeline = CogVideoXPipelineForDDIMInversion.from_pretrained(
-    "THUDM/CogVideoX1.5-5B",
-    torch_dtype=torch.bfloat16,
-).to("cuda")
-
-# Run DDIM inversion, and the videos will be generated in the output_path
-output = pipeline_for_inversion(
-    prompt="prompt that describes the edited video",
-    video_path="path/to/input.mp4",
-    guidance_scale=6.0,
-    num_inference_steps=50,
-    skip_frames_start=0,
-    skip_frames_end=0,
-    frame_sample_step=None,
-    max_num_frames=81,
-    width=720,
-    height=480,
-    seed=42,
-)
-pipeline.export_latents_to_video(output.inverse_latents[-1], "path/to/inverse_video.mp4", fps=8)
-pipeline.export_latents_to_video(output.recon_latents[-1], "path/to/recon_video.mp4", fps=8)
-```

examples/community/cogvideox_ddim_inversion.py

@@ -1,645 +0,0 @@
-"""
-This script performs DDIM inversion for video frames using a pre-trained model and generates
-a video reconstruction based on a provided prompt. It utilizes the CogVideoX pipeline to
-process video frames, apply the DDIM inverse scheduler, and produce an output video.
-
-**Please notice that this script is based on the CogVideoX 5B model, and would not generate
-a good result for 2B variants.**
-
-Usage:
-    python cogvideox_ddim_inversion.py
-        --model-path /path/to/model
-        --prompt "a prompt"
-        --video-path /path/to/video.mp4
-        --output-path /path/to/output
-
-For more details about the cli arguments, please run `python cogvideox_ddim_inversion.py --help`.
-
-Author:
-    LittleNyima <littlenyima[at]163[dot]com>
-"""
-
-import argparse
-import math
-import os
-from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union, cast
-
-import torch
-import torch.nn.functional as F
-import torchvision.transforms as T
-from transformers import T5EncoderModel, T5Tokenizer
-
-from diffusers.models.attention_processor import Attention, CogVideoXAttnProcessor2_0
-from diffusers.models.autoencoders import AutoencoderKLCogVideoX
-from diffusers.models.embeddings import apply_rotary_emb
-from diffusers.models.transformers.cogvideox_transformer_3d import CogVideoXBlock, CogVideoXTransformer3DModel
-from diffusers.pipelines.cogvideo.pipeline_cogvideox import CogVideoXPipeline, retrieve_timesteps
-from diffusers.schedulers import CogVideoXDDIMScheduler, DDIMInverseScheduler
-from diffusers.utils import export_to_video
-
-
-# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error.
-# Very few bug reports but it happens. Look in decord Github issues for more relevant information.
-import decord  # isort: skip
-
-
-class DDIMInversionArguments(TypedDict):
-    model_path: str
-    prompt: str
-    video_path: str
-    output_path: str
-    guidance_scale: float
-    num_inference_steps: int
-    skip_frames_start: int
-    skip_frames_end: int
-    frame_sample_step: Optional[int]
-    max_num_frames: int
-    width: int
-    height: int
-    fps: int
-    dtype: torch.dtype
-    seed: int
-    device: torch.device
-
-
-def get_args() -> DDIMInversionArguments:
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument("--model_path", type=str, required=True, help="Path of the pretrained model")
-    parser.add_argument("--prompt", type=str, required=True, help="Prompt for the direct sample procedure")
-    parser.add_argument("--video_path", type=str, required=True, help="Path of the video for inversion")
-    parser.add_argument("--output_path", type=str, default="output", help="Path of the output videos")
-    parser.add_argument("--guidance_scale", type=float, default=6.0, help="Classifier-free guidance scale")
-    parser.add_argument("--num_inference_steps", type=int, default=50, help="Number of inference steps")
-    parser.add_argument("--skip_frames_start", type=int, default=0, help="Number of skipped frames from the start")
-    parser.add_argument("--skip_frames_end", type=int, default=0, help="Number of skipped frames from the end")
-    parser.add_argument("--frame_sample_step", type=int, default=None, help="Temporal stride of the sampled frames")
-    parser.add_argument("--max_num_frames", type=int, default=81, help="Max number of sampled frames")
-    parser.add_argument("--width", type=int, default=720, help="Resized width of the video frames")
-    parser.add_argument("--height", type=int, default=480, help="Resized height of the video frames")
-    parser.add_argument("--fps", type=int, default=8, help="Frame rate of the output videos")
-    parser.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"], help="Dtype of the model")
-    parser.add_argument("--seed", type=int, default=42, help="Seed for the random number generator")
-    parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"], help="Device for inference")
-
-    args = parser.parse_args()
-    args.dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
-    args.device = torch.device(args.device)
-
-    return DDIMInversionArguments(**vars(args))
-
-
-class CogVideoXAttnProcessor2_0ForDDIMInversion(CogVideoXAttnProcessor2_0):
-    def __init__(self):
-        super().__init__()
-
-    def calculate_attention(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn: Attention,
-        batch_size: int,
-        image_seq_length: int,
-        text_seq_length: int,
-        attention_mask: Optional[torch.Tensor],
-        image_rotary_emb: Optional[torch.Tensor],
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        r"""
-        Core attention computation with inversion-guided RoPE integration.
-
-        Args:
-            query (`torch.Tensor`): `[batch_size, seq_len, dim]` query tensor
-            key (`torch.Tensor`): `[batch_size, seq_len, dim]` key tensor
-            value (`torch.Tensor`): `[batch_size, seq_len, dim]` value tensor
-            attn (`Attention`): Parent attention module with projection layers
-            batch_size (`int`): Effective batch size (after chunk splitting)
-            image_seq_length (`int`): Length of image feature sequence
-            text_seq_length (`int`): Length of text feature sequence
-            attention_mask (`Optional[torch.Tensor]`): Attention mask tensor
-            image_rotary_emb (`Optional[torch.Tensor]`): Rotary embeddings for image positions
-
-        Returns:
-            `Tuple[torch.Tensor, torch.Tensor]`:
-                (1) hidden_states: [batch_size, image_seq_length, dim] processed image features
-                (2) encoder_hidden_states: [batch_size, text_seq_length, dim] processed text features
-        """
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-
-        # Apply RoPE if needed
-        if image_rotary_emb is not None:
-            query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
-            if not attn.is_cross_attention:
-                if key.size(2) == query.size(2):  # Attention for reference hidden states
-                    key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
-                else:  # RoPE should be applied to each group of image tokens
-                    key[:, :, text_seq_length : text_seq_length + image_seq_length] = apply_rotary_emb(
-                        key[:, :, text_seq_length : text_seq_length + image_seq_length], image_rotary_emb
-                    )
-                    key[:, :, text_seq_length * 2 + image_seq_length :] = apply_rotary_emb(
-                        key[:, :, text_seq_length * 2 + image_seq_length :], image_rotary_emb
-                    )
-
-        hidden_states = F.scaled_dot_product_attention(
-            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-        )
-
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        encoder_hidden_states, hidden_states = hidden_states.split(
-            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
-        )
-        return hidden_states, encoder_hidden_states
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        r"""
-        Process the dual-path attention for the inversion-guided denoising procedure.
-
-        Args:
-            attn (`Attention`): Parent attention module
-            hidden_states (`torch.Tensor`): `[batch_size, image_seq_len, dim]` Image tokens
-            encoder_hidden_states (`torch.Tensor`): `[batch_size, text_seq_len, dim]` Text tokens
-            attention_mask (`Optional[torch.Tensor]`): Optional attention mask
-            image_rotary_emb (`Optional[torch.Tensor]`): Rotary embeddings for image tokens
-
-        Returns:
-            `Tuple[torch.Tensor, torch.Tensor]`:
-                (1) Final hidden states: `[batch_size, image_seq_length, dim]` Resulting image tokens
-                (2) Final encoder states: `[batch_size, text_seq_length, dim]` Resulting text tokens
-        """
-        image_seq_length = hidden_states.size(1)
-        text_seq_length = encoder_hidden_states.size(1)
-
-        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
-
-        batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
-        )
-
-        if attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
-
-        query = attn.to_q(hidden_states)
-        key = attn.to_k(hidden_states)
-        value = attn.to_v(hidden_states)
-
-        query, query_reference = query.chunk(2)
-        key, key_reference = key.chunk(2)
-        value, value_reference = value.chunk(2)
-        batch_size = batch_size // 2
-
-        hidden_states, encoder_hidden_states = self.calculate_attention(
-            query=query,
-            key=torch.cat((key, key_reference), dim=1),
-            value=torch.cat((value, value_reference), dim=1),
-            attn=attn,
-            batch_size=batch_size,
-            image_seq_length=image_seq_length,
-            text_seq_length=text_seq_length,
-            attention_mask=attention_mask,
-            image_rotary_emb=image_rotary_emb,
-        )
-        hidden_states_reference, encoder_hidden_states_reference = self.calculate_attention(
-            query=query_reference,
-            key=key_reference,
-            value=value_reference,
-            attn=attn,
-            batch_size=batch_size,
-            image_seq_length=image_seq_length,
-            text_seq_length=text_seq_length,
-            attention_mask=attention_mask,
-            image_rotary_emb=image_rotary_emb,
-        )
-
-        return (
-            torch.cat((hidden_states, hidden_states_reference)),
-            torch.cat((encoder_hidden_states, encoder_hidden_states_reference)),
-        )
-
-
-class OverrideAttnProcessors:
-    r"""
-    Context manager for temporarily overriding attention processors in CogVideo transformer blocks.
-
-    Designed for DDIM inversion process, replaces original attention processors with
-    `CogVideoXAttnProcessor2_0ForDDIMInversion` and restores them upon exit. Uses Python context manager
-    pattern to safely manage processor replacement.
-
-    Typical usage:
-    ```python
-    with OverrideAttnProcessors(transformer):
-        # Perform DDIM inversion operations
-    ```
-
-    Args:
-        transformer (`CogVideoXTransformer3DModel`):
-            The transformer model containing attention blocks to be modified. Should have
-            `transformer_blocks` attribute containing `CogVideoXBlock` instances.
-    """
-
-    def __init__(self, transformer: CogVideoXTransformer3DModel):
-        self.transformer = transformer
-        self.original_processors = {}
-
-    def __enter__(self):
-        for block in self.transformer.transformer_blocks:
-            block = cast(CogVideoXBlock, block)
-            self.original_processors[id(block)] = block.attn1.get_processor()
-            block.attn1.set_processor(CogVideoXAttnProcessor2_0ForDDIMInversion())
-
-    def __exit__(self, _0, _1, _2):
-        for block in self.transformer.transformer_blocks:
-            block = cast(CogVideoXBlock, block)
-            block.attn1.set_processor(self.original_processors[id(block)])
-
-
-def get_video_frames(
-    video_path: str,
-    width: int,
-    height: int,
-    skip_frames_start: int,
-    skip_frames_end: int,
-    max_num_frames: int,
-    frame_sample_step: Optional[int],
-) -> torch.FloatTensor:
-    """
-    Extract and preprocess video frames from a video file for VAE processing.
-
-    Args:
-        video_path (`str`): Path to input video file
-        width (`int`): Target frame width for decoding
-        height (`int`): Target frame height for decoding
-        skip_frames_start (`int`): Number of frames to skip at video start
-        skip_frames_end (`int`): Number of frames to skip at video end
-        max_num_frames (`int`): Maximum allowed number of output frames
-        frame_sample_step (`Optional[int]`):
-            Frame sampling step size. If None, automatically calculated as:
-            (total_frames - skipped_frames) // max_num_frames
-
-    Returns:
-        `torch.FloatTensor`: Preprocessed frames in `[F, C, H, W]` format where:
-        - `F`: Number of frames (adjusted to 4k + 1 for VAE compatibility)
-        - `C`: Channels (3 for RGB)
-        - `H`: Frame height
-        - `W`: Frame width
-    """
-    with decord.bridge.use_torch():
-        video_reader = decord.VideoReader(uri=video_path, width=width, height=height)
-        video_num_frames = len(video_reader)
-        start_frame = min(skip_frames_start, video_num_frames)
-        end_frame = max(0, video_num_frames - skip_frames_end)
-
-        if end_frame <= start_frame:
-            indices = [start_frame]
-        elif end_frame - start_frame <= max_num_frames:
-            indices = list(range(start_frame, end_frame))
-        else:
-            step = frame_sample_step or (end_frame - start_frame) // max_num_frames
-            indices = list(range(start_frame, end_frame, step))
-
-        frames = video_reader.get_batch(indices=indices)
-        frames = frames[:max_num_frames].float()  # ensure that we don't go over the limit
-
-        # Choose first (4k + 1) frames as this is how many is required by the VAE
-        selected_num_frames = frames.size(0)
-        remainder = (3 + selected_num_frames) % 4
-        if remainder != 0:
-            frames = frames[:-remainder]
-        assert frames.size(0) % 4 == 1
-
-        # Normalize the frames
-        transform = T.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)
-        frames = torch.stack(tuple(map(transform, frames)), dim=0)
-
-        return frames.permute(0, 3, 1, 2).contiguous()  # [F, C, H, W]
-
-
-class CogVideoXDDIMInversionOutput:
-    inverse_latents: torch.FloatTensor
-    recon_latents: torch.FloatTensor
-
-    def __init__(self, inverse_latents: torch.FloatTensor, recon_latents: torch.FloatTensor):
-        self.inverse_latents = inverse_latents
-        self.recon_latents = recon_latents
-
-
-class CogVideoXPipelineForDDIMInversion(CogVideoXPipeline):
-    def __init__(
-        self,
-        tokenizer: T5Tokenizer,
-        text_encoder: T5EncoderModel,
-        vae: AutoencoderKLCogVideoX,
-        transformer: CogVideoXTransformer3DModel,
-        scheduler: CogVideoXDDIMScheduler,
-    ):
-        super().__init__(
-            tokenizer=tokenizer,
-            text_encoder=text_encoder,
-            vae=vae,
-            transformer=transformer,
-            scheduler=scheduler,
-        )
-        self.inverse_scheduler = DDIMInverseScheduler(**scheduler.config)
-
-    def encode_video_frames(self, video_frames: torch.FloatTensor) -> torch.FloatTensor:
-        """
-        Encode video frames into latent space using Variational Autoencoder.
-
-        Args:
-            video_frames (`torch.FloatTensor`):
-                Input frames tensor in `[F, C, H, W]` format from `get_video_frames()`
-
-        Returns:
-            `torch.FloatTensor`: Encoded latents in `[1, F, D, H_latent, W_latent]` format where:
-            - `F`: Number of frames (same as input)
-            - `D`: Latent channel dimension
-            - `H_latent`: Latent space height (H // 2^vae.downscale_factor)
-            - `W_latent`: Latent space width (W // 2^vae.downscale_factor)
-        """
-        vae: AutoencoderKLCogVideoX = self.vae
-        video_frames = video_frames.to(device=vae.device, dtype=vae.dtype)
-        video_frames = video_frames.unsqueeze(0).permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
-        latent_dist = vae.encode(x=video_frames).latent_dist.sample().transpose(1, 2)
-        return latent_dist * vae.config.scaling_factor
-
-    @torch.no_grad()
-    def export_latents_to_video(self, latents: torch.FloatTensor, video_path: str, fps: int):
-        r"""
-        Decode latent vectors into video and export as video file.
-
-        Args:
-            latents (`torch.FloatTensor`): Encoded latents in `[B, F, D, H_latent, W_latent]` format from
-                `encode_video_frames()`
-            video_path (`str`): Output path for video file
-            fps (`int`): Target frames per second for output video
-        """
-        video = self.decode_latents(latents)
-        frames = self.video_processor.postprocess_video(video=video, output_type="pil")
-        os.makedirs(os.path.dirname(video_path), exist_ok=True)
-        export_to_video(video_frames=frames[0], output_video_path=video_path, fps=fps)
-
-    # Modified from CogVideoXPipeline.__call__
-    @torch.no_grad()
-    def sample(
-        self,
-        latents: torch.FloatTensor,
-        scheduler: Union[DDIMInverseScheduler, CogVideoXDDIMScheduler],
-        prompt: Optional[Union[str, List[str]]] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 6,
-        use_dynamic_cfg: bool = False,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        reference_latents: torch.FloatTensor = None,
-    ) -> torch.FloatTensor:
-        r"""
-        Execute the core sampling loop for video generation/inversion using CogVideoX.
-
-        Implements the full denoising trajectory recording for both DDIM inversion and
-        generation processes. Supports dynamic classifier-free guidance and reference
-        latent conditioning.
-
-        Args:
-            latents (`torch.FloatTensor`):
-                Initial noise tensor of shape `[B, F, C, H, W]`.
-            scheduler (`Union[DDIMInverseScheduler, CogVideoXDDIMScheduler]`):
-                Scheduling strategy for diffusion process. Use:
-                (1) `DDIMInverseScheduler` for inversion
-                (2) `CogVideoXDDIMScheduler` for generation
-            prompt (`Optional[Union[str, List[str]]]`):
-                Text prompt(s) for conditional generation. Defaults to unconditional.
-            negative_prompt (`Optional[Union[str, List[str]]]`):
-                Negative prompt(s) for guidance. Requires `guidance_scale > 1`.
-            num_inference_steps (`int`):
-                Number of denoising steps. Affects quality/compute trade-off.
-            guidance_scale (`float`):
-                Classifier-free guidance weight. 1.0 = no guidance.
-            use_dynamic_cfg (`bool`):
-                Enable time-varying guidance scale (cosine schedule)
-            eta (`float`):
-                DDIM variance parameter (0 = deterministic process)
-            generator (`Optional[Union[torch.Generator, List[torch.Generator]]]`):
-                Random number generator(s) for reproducibility
-            attention_kwargs (`Optional[Dict[str, Any]]`):
-                Custom parameters for attention modules
-            reference_latents (`torch.FloatTensor`):
-                Reference latent trajectory for conditional sampling. Shape should match
-                `[T, B, F, C, H, W]` where `T` is number of timesteps
-
-        Returns:
-            `torch.FloatTensor`:
-                Full denoising trajectory tensor of shape `[T, B, F, C, H, W]`.
-        """
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-
-        device = self._execution_device
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 3. Encode input prompt
-        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-            prompt,
-            negative_prompt,
-            do_classifier_free_guidance,
-            device=device,
-        )
-        if do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-        if reference_latents is not None:
-            prompt_embeds = torch.cat([prompt_embeds] * 2, dim=0)
-
-        # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps, device)
-        self._num_timesteps = len(timesteps)
-
-        # 5. Prepare latents.
-        latents = latents.to(device=device) * scheduler.init_noise_sigma
-
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-        if isinstance(scheduler, DDIMInverseScheduler):  # Inverse scheduler does not accept extra kwargs
-            extra_step_kwargs = {}
-
-        # 7. Create rotary embeds if required
-        image_rotary_emb = (
-            self._prepare_rotary_positional_embeddings(
-                height=latents.size(3) * self.vae_scale_factor_spatial,
-                width=latents.size(4) * self.vae_scale_factor_spatial,
-                num_frames=latents.size(1),
-                device=device,
-            )
-            if self.transformer.config.use_rotary_positional_embeddings
-            else None
-        )
-
-        # 8. Denoising loop
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * scheduler.order, 0)
-
-        trajectory = torch.zeros_like(latents).unsqueeze(0).repeat(len(timesteps), 1, 1, 1, 1, 1)
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                if reference_latents is not None:
-                    reference = reference_latents[i]
-                    reference = torch.cat([reference] * 2) if do_classifier_free_guidance else reference
-                    latent_model_input = torch.cat([latent_model_input, reference], dim=0)
-                latent_model_input = scheduler.scale_model_input(latent_model_input, t)
-
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0])
-
-                # predict noise model_output
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    encoder_hidden_states=prompt_embeds,
-                    timestep=timestep,
-                    image_rotary_emb=image_rotary_emb,
-                    attention_kwargs=attention_kwargs,
-                    return_dict=False,
-                )[0]
-                noise_pred = noise_pred.float()
-
-                if reference_latents is not None:  # Recover the original batch size
-                    noise_pred, _ = noise_pred.chunk(2)
-
-                # perform guidance
-                if use_dynamic_cfg:
-                    self._guidance_scale = 1 + guidance_scale * (
-                        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
-                    )
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the noisy sample x_t-1 -> x_t
-                latents = scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-                latents = latents.to(prompt_embeds.dtype)
-                trajectory[i] = latents
-
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0):
-                    progress_bar.update()
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        return trajectory
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt: str,
-        video_path: str,
-        guidance_scale: float,
-        num_inference_steps: int,
-        skip_frames_start: int,
-        skip_frames_end: int,
-        frame_sample_step: Optional[int],
-        max_num_frames: int,
-        width: int,
-        height: int,
-        seed: int,
-    ):
-        """
-        Performs DDIM inversion on a video to reconstruct it with a new prompt.
-
-        Args:
-            prompt (`str`): The text prompt to guide the reconstruction.
-            video_path (`str`): Path to the input video file.
-            guidance_scale (`float`): Scale for classifier-free guidance.
-            num_inference_steps (`int`): Number of denoising steps.
-            skip_frames_start (`int`): Number of frames to skip from the beginning of the video.
-            skip_frames_end (`int`): Number of frames to skip from the end of the video.
-            frame_sample_step (`Optional[int]`): Step size for sampling frames. If None, all frames are used.
-            max_num_frames (`int`): Maximum number of frames to process.
-            width (`int`): Width of the output video frames.
-            height (`int`): Height of the output video frames.
-            seed (`int`): Random seed for reproducibility.
-
-        Returns:
-            `CogVideoXDDIMInversionOutput`: Contains the inverse latents and reconstructed latents.
-        """
-        if not self.transformer.config.use_rotary_positional_embeddings:
-            raise NotImplementedError("This script supports CogVideoX 5B model only.")
-        video_frames = get_video_frames(
-            video_path=video_path,
-            width=width,
-            height=height,
-            skip_frames_start=skip_frames_start,
-            skip_frames_end=skip_frames_end,
-            max_num_frames=max_num_frames,
-            frame_sample_step=frame_sample_step,
-        ).to(device=self.device)
-        video_latents = self.encode_video_frames(video_frames=video_frames)
-        inverse_latents = self.sample(
-            latents=video_latents,
-            scheduler=self.inverse_scheduler,
-            prompt="",
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            generator=torch.Generator(device=self.device).manual_seed(seed),
-        )
-        with OverrideAttnProcessors(transformer=self.transformer):
-            recon_latents = self.sample(
-                latents=torch.randn_like(video_latents),
-                scheduler=self.scheduler,
-                prompt=prompt,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                generator=torch.Generator(device=self.device).manual_seed(seed),
-                reference_latents=reversed(inverse_latents),
-            )
-        return CogVideoXDDIMInversionOutput(
-            inverse_latents=inverse_latents,
-            recon_latents=recon_latents,
-        )
-
-
-if __name__ == "__main__":
-    arguments = get_args()
-    pipeline = CogVideoXPipelineForDDIMInversion.from_pretrained(
-        arguments.pop("model_path"),
-        torch_dtype=arguments.pop("dtype"),
-    ).to(device=arguments.pop("device"))
-
-    output_path = arguments.pop("output_path")
-    fps = arguments.pop("fps")
-    inverse_video_path = os.path.join(output_path, f"{arguments.get('video_path')}_inversion.mp4")
-    recon_video_path = os.path.join(output_path, f"{arguments.get('video_path')}_reconstruction.mp4")
-
-    # Run DDIM inversion
-    output = pipeline(**arguments)
-    pipeline.export_latents_to_video(output.inverse_latents[-1], inverse_video_path, fps)
-    pipeline.export_latents_to_video(output.recon_latents[-1], recon_video_path, fps)

examples/research_projects/realfill/requirements.txt

@@ -6,4 +6,4 @@ torch==2.2.0
 torchvision>=0.16
 ftfy==6.1.1
 tensorboard==2.14.0
-Jinja2==3.1.6
+Jinja2==3.1.5

setup.py

@@ -128,10 +128,6 @@ _deps = [
     "GitPython<3.1.19",
     "scipy",
     "onnx",
-    "optimum_quanto>=0.2.6",
-    "gguf>=0.10.0",
-    "torchao>=0.7.0",
-    "bitsandbytes>=0.43.3",
     "regex!=2019.12.17",
     "requests",
     "tensorboard",
@@ -239,11 +235,6 @@ extras["test"] = deps_list(
 )
 extras["torch"] = deps_list("torch", "accelerate")
 
-extras["bitsandbytes"] = deps_list("bitsandbytes", "accelerate")
-extras["gguf"] = deps_list("gguf", "accelerate")
-extras["optimum_quanto"] = deps_list("optimum_quanto", "accelerate")
-extras["torchao"] = deps_list("torchao", "accelerate")
-
 if os.name == "nt":  # windows
     extras["flax"] = []  # jax is not supported on windows
 else:

src/diffusers/__init__.py

@@ -2,15 +2,6 @@ __version__ = "0.33.0.dev0"
 
 from typing import TYPE_CHECKING
 
-from diffusers.quantizers import quantization_config
-from diffusers.utils import dummy_gguf_objects
-from diffusers.utils.import_utils import (
-    is_bitsandbytes_available,
-    is_gguf_available,
-    is_optimum_quanto_version,
-    is_torchao_available,
-)
-
 from .utils import (
     DIFFUSERS_SLOW_IMPORT,
     OptionalDependencyNotAvailable,
@@ -20,7 +11,6 @@ from .utils import (
     is_librosa_available,
     is_note_seq_available,
     is_onnx_available,
-    is_optimum_quanto_available,
     is_scipy_available,
     is_sentencepiece_available,
     is_torch_available,
@@ -42,7 +32,7 @@ _import_structure = {
     "loaders": ["FromOriginalModelMixin"],
     "models": [],
     "pipelines": [],
-    "quantizers.quantization_config": [],
+    "quantizers.quantization_config": ["BitsAndBytesConfig", "GGUFQuantizationConfig", "TorchAoConfig"],
     "schedulers": [],
     "utils": [
         "OptionalDependencyNotAvailable",
@@ -64,55 +54,6 @@ _import_structure = {
     ],
 }
 
-try:
-    if not is_bitsandbytes_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_bitsandbytes_objects
-
-    _import_structure["utils.dummy_bitsandbytes_objects"] = [
-        name for name in dir(dummy_bitsandbytes_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("BitsAndBytesConfig")
-
-try:
-    if not is_gguf_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_gguf_objects
-
-    _import_structure["utils.dummy_gguf_objects"] = [
-        name for name in dir(dummy_gguf_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig")
-
-try:
-    if not is_torchao_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torchao_objects
-
-    _import_structure["utils.dummy_torchao_objects"] = [
-        name for name in dir(dummy_torchao_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("TorchAoConfig")
-
-try:
-    if not is_optimum_quanto_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_optimum_quanto_objects
-
-    _import_structure["utils.dummy_optimum_quanto_objects"] = [
-        name for name in dir(dummy_optimum_quanto_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("QuantoConfig")
-
-
 try:
     if not is_onnx_available():
         raise OptionalDependencyNotAvailable()
@@ -657,38 +598,7 @@ else:
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     from .configuration_utils import ConfigMixin
-
-    try:
-        if not is_bitsandbytes_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_bitsandbytes_objects import *
-    else:
-        from .quantizers.quantization_config import BitsAndBytesConfig
-
-    try:
-        if not is_gguf_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_gguf_objects import *
-    else:
-        from .quantizers.quantization_config import GGUFQuantizationConfig
-
-    try:
-        if not is_torchao_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torchao_objects import *
-    else:
-        from .quantizers.quantization_config import TorchAoConfig
-
-    try:
-        if not is_optimum_quanto_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_optimum_quanto_objects import *
-    else:
-        from .quantizers.quantization_config import QuantoConfig
+    from .quantizers.quantization_config import BitsAndBytesConfig, GGUFQuantizationConfig, TorchAoConfig
 
     try:
         if not is_onnx_available():

src/diffusers/dependency_versions_table.py

@@ -35,10 +35,6 @@ deps = {
     "GitPython": "GitPython<3.1.19",
     "scipy": "scipy",
     "onnx": "onnx",
-    "optimum_quanto": "optimum_quanto>=0.2.6",
-    "gguf": "gguf>=0.10.0",
-    "torchao": "torchao>=0.7.0",
-    "bitsandbytes": "bitsandbytes>=0.43.3",
     "regex": "regex!=2019.12.17",
     "requests": "requests",
     "tensorboard": "tensorboard",

src/diffusers/loaders/lora_conversion_utils.py

@@ -654,7 +654,6 @@ def _convert_kohya_flux_lora_to_diffusers(state_dict):
 
                 _convert(k, diffusers_key, state_dict, new_state_dict)
 
-        remaining_all_unet = False
         if state_dict:
             remaining_all_unet = all(k.startswith("lora_unet_") for k in state_dict)
         if remaining_all_unet:

src/diffusers/loaders/single_file_utils.py

@@ -397,7 +397,6 @@ def load_single_file_checkpoint(
 
     else:
         repo_id, weights_name = _extract_repo_id_and_weights_name(pretrained_model_link_or_path)
-        user_agent = {"file_type": "single_file", "framework": "pytorch"}
         pretrained_model_link_or_path = _get_model_file(
             repo_id,
             weights_name=weights_name,
@@ -407,7 +406,6 @@ def load_single_file_checkpoint(
             local_files_only=local_files_only,
             token=token,
             revision=revision,
-            user_agent=user_agent,
         )
 
     checkpoint = load_state_dict(pretrained_model_link_or_path, disable_mmap=disable_mmap)

src/diffusers/models/model_loading_utils.py

@@ -245,9 +245,6 @@ def load_model_dict_into_meta(
             ):
                 param = param.to(torch.float32)
                 set_module_kwargs["dtype"] = torch.float32
-            # For quantizers have save weights using torch.float8_e4m3fn
-            elif hf_quantizer is not None and param.dtype == getattr(torch, "float8_e4m3fn", None):
-                pass
             else:
                 param = param.to(dtype)
                 set_module_kwargs["dtype"] = dtype
@@ -295,9 +292,7 @@ def load_model_dict_into_meta(
         elif is_quantized and (
             hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict, param_device=param_device)
         ):
-            hf_quantizer.create_quantized_param(
-                model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype
-            )
+            hf_quantizer.create_quantized_param(model, param, param_name, param_device, state_dict, unexpected_keys)
         else:
             set_module_tensor_to_device(model, param_name, param_device, value=param, **set_module_kwargs)
 

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -694,7 +694,7 @@ class FluxPipeline(
                 Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
                 their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
                 will be used.
-            guidance_scale (`float`, *optional*, defaults to 3.5):
+            guidance_scale (`float`, *optional*, defaults to 7.0):
                 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                 `guidance_scale` is defined as `w` of equation 2. of [Imagen
                 Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >

src/diffusers/pipelines/flux/pipeline_flux_control.py

@@ -660,7 +660,7 @@ class FluxControlPipeline(
                 Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
                 their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
                 will be used.
-            guidance_scale (`float`, *optional*, defaults to 3.5):
+            guidance_scale (`float`, *optional*, defaults to 7.0):
                 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                 `guidance_scale` is defined as `w` of equation 2. of [Imagen
                 Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >

src/diffusers/pipelines/flux/pipeline_flux_controlnet.py

@@ -202,7 +202,7 @@ class FluxControlNetPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleF
 
     model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae"
     _optional_components = ["image_encoder", "feature_extractor"]
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "control_image"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
 
     def __init__(
         self,
@@ -1149,7 +1149,6 @@ class FluxControlNetPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleF
 
                     latents = callback_outputs.pop("latents", latents)
                     prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    control_image = callback_outputs.pop("control_image", control_image)
 
                 # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):

src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py

@@ -198,7 +198,7 @@ class FluxControlNetImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, From
 
     model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
     _optional_components = []
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "control_image"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
 
     def __init__(
         self,
@@ -973,7 +973,6 @@ class FluxControlNetImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin, From
 
                     latents = callback_outputs.pop("latents", latents)
                     prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    control_image = callback_outputs.pop("control_image", control_image)
 
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()

src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py

@@ -200,7 +200,7 @@ class FluxControlNetInpaintPipeline(DiffusionPipeline, FluxLoraLoaderMixin, From
 
     model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
     _optional_components = []
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "control_image", "mask", "masked_image_latents"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
 
     def __init__(
         self,
@@ -1178,9 +1178,6 @@ class FluxControlNetInpaintPipeline(DiffusionPipeline, FluxLoraLoaderMixin, From
 
                     latents = callback_outputs.pop("latents", latents)
                     prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    control_image = callback_outputs.pop("control_image", control_image)
-                    mask = callback_outputs.pop("mask", mask)
-                    masked_image_latents = callback_outputs.pop("masked_image_latents", masked_image_latents)
 
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()

src/diffusers/pipelines/flux/pipeline_flux_fill.py

@@ -738,7 +738,7 @@ class FluxFillPipeline(
                 Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
                 their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
                 will be used.
-            guidance_scale (`float`, *optional*, defaults to 30.0):
+            guidance_scale (`float`, *optional*, defaults to 7.0):
                 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                 `guidance_scale` is defined as `w` of equation 2. of [Imagen
                 Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >

src/diffusers/quantizers/auto.py

@@ -26,10 +26,8 @@ from .quantization_config import (
     GGUFQuantizationConfig,
     QuantizationConfigMixin,
     QuantizationMethod,
-    QuantoConfig,
     TorchAoConfig,
 )
-from .quanto import QuantoQuantizer
 from .torchao import TorchAoHfQuantizer
 
 
@@ -37,7 +35,6 @@ AUTO_QUANTIZER_MAPPING = {
     "bitsandbytes_4bit": BnB4BitDiffusersQuantizer,
     "bitsandbytes_8bit": BnB8BitDiffusersQuantizer,
     "gguf": GGUFQuantizer,
-    "quanto": QuantoQuantizer,
     "torchao": TorchAoHfQuantizer,
 }
 
@@ -45,7 +42,6 @@ AUTO_QUANTIZATION_CONFIG_MAPPING = {
     "bitsandbytes_4bit": BitsAndBytesConfig,
     "bitsandbytes_8bit": BitsAndBytesConfig,
     "gguf": GGUFQuantizationConfig,
-    "quanto": QuantoConfig,
     "torchao": TorchAoConfig,
 }
 

src/diffusers/quantizers/quantization_config.py

@@ -45,7 +45,6 @@ class QuantizationMethod(str, Enum):
     BITS_AND_BYTES = "bitsandbytes"
     GGUF = "gguf"
     TORCHAO = "torchao"
-    QUANTO = "quanto"
 
 
 if is_torchao_available():
@@ -687,38 +686,3 @@ class TorchAoConfig(QuantizationConfigMixin):
         return (
             f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True, cls=TorchAoJSONEncoder)}\n"
         )
-
-
-@dataclass
-class QuantoConfig(QuantizationConfigMixin):
-    """
-    This is a wrapper class about all possible attributes and features that you can play with a model that has been
-    loaded using `quanto`.
-
-    Args:
-        weights_dtype (`str`, *optional*, defaults to `"int8"`):
-            The target dtype for the weights after quantization. Supported values are ("float8","int8","int4","int2")
-       modules_to_not_convert (`list`, *optional*, default to `None`):
-            The list of modules to not quantize, useful for quantizing models that explicitly require to have some
-            modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
-    """
-
-    def __init__(
-        self,
-        weights_dtype: str = "int8",
-        modules_to_not_convert: Optional[List[str]] = None,
-        **kwargs,
-    ):
-        self.quant_method = QuantizationMethod.QUANTO
-        self.weights_dtype = weights_dtype
-        self.modules_to_not_convert = modules_to_not_convert
-
-        self.post_init()
-
-    def post_init(self):
-        r"""
-        Safety checker that arguments are correct
-        """
-        accepted_weights = ["float8", "int8", "int4", "int2"]
-        if self.weights_dtype not in accepted_weights:
-            raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")

src/diffusers/quantizers/quanto/__init__.py

@@ -1 +0,0 @@
-from .quanto_quantizer import QuantoQuantizer

src/diffusers/quantizers/quanto/quanto_quantizer.py

@@ -1,177 +0,0 @@
-from typing import TYPE_CHECKING, Any, Dict, List, Union
-
-from diffusers.utils.import_utils import is_optimum_quanto_version
-
-from ...utils import (
-    get_module_from_name,
-    is_accelerate_available,
-    is_accelerate_version,
-    is_optimum_quanto_available,
-    is_torch_available,
-    logging,
-)
-from ..base import DiffusersQuantizer
-
-
-if TYPE_CHECKING:
-    from ...models.modeling_utils import ModelMixin
-
-
-if is_torch_available():
-    import torch
-
-if is_accelerate_available():
-    from accelerate.utils import CustomDtype, set_module_tensor_to_device
-
-if is_optimum_quanto_available():
-    from .utils import _replace_with_quanto_layers
-
-logger = logging.get_logger(__name__)
-
-
-class QuantoQuantizer(DiffusersQuantizer):
-    r"""
-    Diffusers Quantizer for Optimum Quanto
-    """
-
-    use_keep_in_fp32_modules = True
-    requires_calibration = False
-    required_packages = ["quanto", "accelerate"]
-
-    def __init__(self, quantization_config, **kwargs):
-        super().__init__(quantization_config, **kwargs)
-
-    def validate_environment(self, *args, **kwargs):
-        if not is_optimum_quanto_available():
-            raise ImportError(
-                "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)"
-            )
-        if not is_optimum_quanto_version(">=", "0.2.6"):
-            raise ImportError(
-                "Loading an optimum-quanto quantized model requires `optimum-quanto>=0.2.6`. "
-                "Please upgrade your installation with `pip install --upgrade optimum-quanto"
-            )
-
-        if not is_accelerate_available():
-            raise ImportError(
-                "Loading an optimum-quanto quantized model requires accelerate library (`pip install accelerate`)"
-            )
-
-        device_map = kwargs.get("device_map", None)
-        if isinstance(device_map, dict) and len(device_map.keys()) > 1:
-            raise ValueError(
-                "`device_map` for multi-GPU inference or CPU/disk offload is currently not supported with Diffusers and the Quanto backend"
-            )
-
-    def check_if_quantized_param(
-        self,
-        model: "ModelMixin",
-        param_value: "torch.Tensor",
-        param_name: str,
-        state_dict: Dict[str, Any],
-        **kwargs,
-    ):
-        # Quanto imports diffusers internally. This is here to prevent circular imports
-        from optimum.quanto import QModuleMixin, QTensor
-        from optimum.quanto.tensor.packed import PackedTensor
-
-        module, tensor_name = get_module_from_name(model, param_name)
-        if self.pre_quantized and any(isinstance(module, t) for t in [QTensor, PackedTensor]):
-            return True
-        elif isinstance(module, QModuleMixin) and "weight" in tensor_name:
-            return not module.frozen
-
-        return False
-
-    def create_quantized_param(
-        self,
-        model: "ModelMixin",
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        *args,
-        **kwargs,
-    ):
-        """
-        Create the quantized parameter by calling .freeze() after setting it to the module.
-        """
-
-        dtype = kwargs.get("dtype", torch.float32)
-        module, tensor_name = get_module_from_name(model, param_name)
-        if self.pre_quantized:
-            setattr(module, tensor_name, param_value)
-        else:
-            set_module_tensor_to_device(model, param_name, target_device, param_value, dtype)
-            module.freeze()
-            module.weight.requires_grad = False
-
-    def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
-        max_memory = {key: val * 0.90 for key, val in max_memory.items()}
-        return max_memory
-
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
-        if is_accelerate_version(">=", "0.27.0"):
-            mapping = {
-                "int8": torch.int8,
-                "float8": CustomDtype.FP8,
-                "int4": CustomDtype.INT4,
-                "int2": CustomDtype.INT2,
-            }
-            target_dtype = mapping[self.quantization_config.weights_dtype]
-
-        return target_dtype
-
-    def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype":
-        if torch_dtype is None:
-            logger.info("You did not specify `torch_dtype` in `from_pretrained`. Setting it to `torch.float32`.")
-            torch_dtype = torch.float32
-        return torch_dtype
-
-    def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]:
-        # Quanto imports diffusers internally. This is here to prevent circular imports
-        from optimum.quanto import QModuleMixin
-
-        not_missing_keys = []
-        for name, module in model.named_modules():
-            if isinstance(module, QModuleMixin):
-                for missing in missing_keys:
-                    if (
-                        (name in missing or name in f"{prefix}.{missing}")
-                        and not missing.endswith(".weight")
-                        and not missing.endswith(".bias")
-                    ):
-                        not_missing_keys.append(missing)
-        return [k for k in missing_keys if k not in not_missing_keys]
-
-    def _process_model_before_weight_loading(
-        self,
-        model: "ModelMixin",
-        device_map,
-        keep_in_fp32_modules: List[str] = [],
-        **kwargs,
-    ):
-        self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
-
-        if not isinstance(self.modules_to_not_convert, list):
-            self.modules_to_not_convert = [self.modules_to_not_convert]
-
-        self.modules_to_not_convert.extend(keep_in_fp32_modules)
-
-        model = _replace_with_quanto_layers(
-            model,
-            modules_to_not_convert=self.modules_to_not_convert,
-            quantization_config=self.quantization_config,
-            pre_quantized=self.pre_quantized,
-        )
-        model.config.quantization_config = self.quantization_config
-
-    def _process_model_after_weight_loading(self, model, **kwargs):
-        return model
-
-    @property
-    def is_trainable(self):
-        return True
-
-    @property
-    def is_serializable(self):
-        return True

src/diffusers/quantizers/quanto/utils.py

@@ -1,60 +0,0 @@
-import torch.nn as nn
-
-from ...utils import is_accelerate_available, logging
-
-
-logger = logging.get_logger(__name__)
-
-if is_accelerate_available():
-    from accelerate import init_empty_weights
-
-
-def _replace_with_quanto_layers(model, quantization_config, modules_to_not_convert: list, pre_quantized=False):
-    # Quanto imports diffusers internally. These are placed here to avoid circular imports
-    from optimum.quanto import QLinear, freeze, qfloat8, qint2, qint4, qint8
-
-    def _get_weight_type(dtype: str):
-        return {"float8": qfloat8, "int8": qint8, "int4": qint4, "int2": qint2}[dtype]
-
-    def _replace_layers(model, quantization_config, modules_to_not_convert):
-        has_children = list(model.children())
-        if not has_children:
-            return model
-
-        for name, module in model.named_children():
-            _replace_layers(module, quantization_config, modules_to_not_convert)
-
-            if name in modules_to_not_convert:
-                continue
-
-            if isinstance(module, nn.Linear):
-                with init_empty_weights():
-                    qlinear = QLinear(
-                        in_features=module.in_features,
-                        out_features=module.out_features,
-                        bias=module.bias is not None,
-                        dtype=module.weight.dtype,
-                        weights=_get_weight_type(quantization_config.weights_dtype),
-                    )
-                    model._modules[name] = qlinear
-                    model._modules[name].source_cls = type(module)
-                    model._modules[name].requires_grad_(False)
-
-        return model
-
-    model = _replace_layers(model, quantization_config, modules_to_not_convert)
-    has_been_replaced = any(isinstance(replaced_module, QLinear) for _, replaced_module in model.named_modules())
-
-    if not has_been_replaced:
-        logger.warning(
-            f"{model.__class__.__name__} does not appear to have any `nn.Linear` modules. Quantization will not be applied."
-            " Please check your model architecture, or submit an issue on Github if you think this is a bug."
-            " https://github.com/huggingface/diffusers/issues/new"
-        )
-
-    # We need to freeze the pre_quantized model in order for the loaded state_dict and model state dict
-    # to match when trying to load weights with load_model_dict_into_meta
-    if pre_quantized:
-        freeze(model)
-
-    return model

src/diffusers/utils/__init__.py

@@ -79,8 +79,6 @@ from .import_utils import (
     is_matplotlib_available,
     is_note_seq_available,
     is_onnx_available,
-    is_optimum_quanto_available,
-    is_optimum_quanto_version,
     is_peft_available,
     is_peft_version,
     is_safetensors_available,

src/diffusers/utils/dummy_bitsandbytes_objects.py

@@ -1,17 +0,0 @@
-# This file is autogenerated by the command `make fix-copies`, do not edit.
-from ..utils import DummyObject, requires_backends
-
-
-class BitsAndBytesConfig(metaclass=DummyObject):
-    _backends = ["bitsandbytes"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["bitsandbytes"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["bitsandbytes"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["bitsandbytes"])

src/diffusers/utils/dummy_gguf_objects.py

@@ -1,17 +0,0 @@
-# This file is autogenerated by the command `make fix-copies`, do not edit.
-from ..utils import DummyObject, requires_backends
-
-
-class GGUFQuantizationConfig(metaclass=DummyObject):
-    _backends = ["gguf"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["gguf"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["gguf"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["gguf"])

src/diffusers/utils/dummy_optimum_quanto_objects.py

@@ -1,17 +0,0 @@
-# This file is autogenerated by the command `make fix-copies`, do not edit.
-from ..utils import DummyObject, requires_backends
-
-
-class QuantoConfig(metaclass=DummyObject):
-    _backends = ["optimum_quanto"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["optimum_quanto"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["optimum_quanto"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["optimum_quanto"])

src/diffusers/utils/dummy_torchao_objects.py

@@ -1,17 +0,0 @@
-# This file is autogenerated by the command `make fix-copies`, do not edit.
-from ..utils import DummyObject, requires_backends
-
-
-class TorchAoConfig(metaclass=DummyObject):
-    _backends = ["torchao"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["torchao"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["torchao"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["torchao"])

src/diffusers/utils/import_utils.py

@@ -365,15 +365,6 @@ if _is_torchao_available:
         _is_torchao_available = False
 
 
-_is_optimum_quanto_available = importlib.util.find_spec("optimum") is not None
-if _is_optimum_quanto_available:
-    try:
-        _optimum_quanto_version = importlib_metadata.version("optimum_quanto")
-        logger.debug(f"Successfully import optimum-quanto version {_optimum_quanto_version}")
-    except importlib_metadata.PackageNotFoundError:
-        _is_optimum_quanto_available = False
-
-
 def is_torch_available():
     return _torch_available
 
@@ -502,10 +493,6 @@ def is_torchao_available():
     return _is_torchao_available
 
 
-def is_optimum_quanto_available():
-    return _is_optimum_quanto_available
-
-
 # docstyle-ignore
 FLAX_IMPORT_ERROR = """
 {0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the
@@ -649,11 +636,6 @@ TORCHAO_IMPORT_ERROR = """
 torchao`
 """
 
-QUANTO_IMPORT_ERROR = """
-{0} requires the optimum-quanto library but it was not found in your environment. You can install it with pip: `pip
-install optimum-quanto`
-"""
-
 BACKENDS_MAPPING = OrderedDict(
     [
         ("bs4", (is_bs4_available, BS4_IMPORT_ERROR)),
@@ -681,7 +663,6 @@ BACKENDS_MAPPING = OrderedDict(
         ("imageio", (is_imageio_available, IMAGEIO_IMPORT_ERROR)),
         ("gguf", (is_gguf_available, GGUF_IMPORT_ERROR)),
         ("torchao", (is_torchao_available, TORCHAO_IMPORT_ERROR)),
-        ("quanto", (is_optimum_quanto_available, QUANTO_IMPORT_ERROR)),
     ]
 )
 
@@ -883,21 +864,6 @@ def is_k_diffusion_version(operation: str, version: str):
     return compare_versions(parse(_k_diffusion_version), operation, version)
 
 
-def is_optimum_quanto_version(operation: str, version: str):
-    """
-    Compares the current Accelerate version to a given reference with an operation.
-
-    Args:
-        operation (`str`):
-            A string representation of an operator, such as `">"` or `"<="`
-        version (`str`):
-            A version string
-    """
-    if not _is_optimum_quanto_available:
-        return False
-    return compare_versions(parse(_optimum_quanto_version), operation, version)
-
-
 def get_objects_from_module(module):
     """
     Returns a dict of object names and values in a module, while skipping private/internal objects

tests/quantization/quanto/test_quanto.py

@@ -1,346 +0,0 @@
-import gc
-import tempfile
-import unittest
-
-from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig
-from diffusers.models.attention_processor import Attention
-from diffusers.utils import is_optimum_quanto_available, is_torch_available
-from diffusers.utils.testing_utils import (
-    nightly,
-    numpy_cosine_similarity_distance,
-    require_accelerate,
-    require_big_gpu_with_torch_cuda,
-    torch_device,
-)
-
-
-if is_optimum_quanto_available():
-    from optimum.quanto import QLinear
-
-if is_torch_available():
-    import torch
-    import torch.nn as nn
-
-    class LoRALayer(nn.Module):
-        """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only
-
-        Taken from
-        https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_4bit.py#L62C5-L78C77
-        """
-
-        def __init__(self, module: nn.Module, rank: int):
-            super().__init__()
-            self.module = module
-            self.adapter = nn.Sequential(
-                nn.Linear(module.in_features, rank, bias=False),
-                nn.Linear(rank, module.out_features, bias=False),
-            )
-            small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5
-            nn.init.normal_(self.adapter[0].weight, std=small_std)
-            nn.init.zeros_(self.adapter[1].weight)
-            self.adapter.to(module.weight.device)
-
-        def forward(self, input, *args, **kwargs):
-            return self.module(input, *args, **kwargs) + self.adapter(input)
-
-
-@nightly
-@require_big_gpu_with_torch_cuda
-@require_accelerate
-class QuantoBaseTesterMixin:
-    model_id = None
-    pipeline_model_id = None
-    model_cls = None
-    torch_dtype = torch.bfloat16
-    # the expected reduction in peak memory used compared to an unquantized model expressed as a percentage
-    expected_memory_reduction = 0.0
-    keep_in_fp32_module = ""
-    modules_to_not_convert = ""
-    _test_torch_compile = False
-
-    def setUp(self):
-        torch.cuda.reset_peak_memory_stats()
-        torch.cuda.empty_cache()
-        gc.collect()
-
-    def tearDown(self):
-        torch.cuda.reset_peak_memory_stats()
-        torch.cuda.empty_cache()
-        gc.collect()
-
-    def get_dummy_init_kwargs(self):
-        return {"weights_dtype": "float8"}
-
-    def get_dummy_model_init_kwargs(self):
-        return {
-            "pretrained_model_name_or_path": self.model_id,
-            "torch_dtype": self.torch_dtype,
-            "quantization_config": QuantoConfig(**self.get_dummy_init_kwargs()),
-        }
-
-    def test_quanto_layers(self):
-        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
-        for name, module in model.named_modules():
-            if isinstance(module, torch.nn.Linear):
-                assert isinstance(module, QLinear)
-
-    def test_quanto_memory_usage(self):
-        unquantized_model = self.model_cls.from_pretrained(self.model_id, torch_dtype=self.torch_dtype)
-        unquantized_model_memory = unquantized_model.get_memory_footprint() / 1024**3
-
-        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
-        inputs = self.get_dummy_inputs()
-
-        torch.cuda.reset_peak_memory_stats()
-        torch.cuda.empty_cache()
-
-        model.to(torch_device)
-        with torch.no_grad():
-            model(**inputs)
-        max_memory = torch.cuda.max_memory_allocated() / 1024**3
-        assert (1.0 - (max_memory / unquantized_model_memory)) >= self.expected_memory_reduction
-
-    def test_keep_modules_in_fp32(self):
-        r"""
-        A simple tests to check if the modules under `_keep_in_fp32_modules` are kept in fp32.
-        Also ensures if inference works.
-        """
-        _keep_in_fp32_modules = self.model_cls._keep_in_fp32_modules
-        self.model_cls._keep_in_fp32_modules = self.keep_in_fp32_module
-
-        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
-        model.to("cuda")
-
-        for name, module in model.named_modules():
-            if isinstance(module, torch.nn.Linear):
-                if name in model._keep_in_fp32_modules:
-                    assert module.weight.dtype == torch.float32
-        self.model_cls._keep_in_fp32_modules = _keep_in_fp32_modules
-
-    def test_modules_to_not_convert(self):
-        init_kwargs = self.get_dummy_model_init_kwargs()
-
-        quantization_config_kwargs = self.get_dummy_init_kwargs()
-        quantization_config_kwargs.update({"modules_to_not_convert": self.modules_to_not_convert})
-        quantization_config = QuantoConfig(**quantization_config_kwargs)
-
-        init_kwargs.update({"quantization_config": quantization_config})
-
-        model = self.model_cls.from_pretrained(**init_kwargs)
-        model.to("cuda")
-
-        for name, module in model.named_modules():
-            if name in self.modules_to_not_convert:
-                assert not isinstance(module, QLinear)
-
-    def test_dtype_assignment(self):
-        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
-
-        with self.assertRaises(ValueError):
-            # Tries with a `dtype`
-            model.to(torch.float16)
-
-        with self.assertRaises(ValueError):
-            # Tries with a `device` and `dtype`
-            model.to(device="cuda:0", dtype=torch.float16)
-
-        with self.assertRaises(ValueError):
-            # Tries with a cast
-            model.float()
-
-        with self.assertRaises(ValueError):
-            # Tries with a cast
-            model.half()
-
-        # This should work
-        model.to("cuda")
-
-    def test_serialization(self):
-        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
-        inputs = self.get_dummy_inputs()
-
-        model.to(torch_device)
-        with torch.no_grad():
-            model_output = model(**inputs)
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            model.save_pretrained(tmp_dir)
-            saved_model = self.model_cls.from_pretrained(
-                tmp_dir,
-                torch_dtype=torch.bfloat16,
-            )
-
-        saved_model.to(torch_device)
-        with torch.no_grad():
-            saved_model_output = saved_model(**inputs)
-
-        assert torch.allclose(model_output.sample, saved_model_output.sample, rtol=1e-5, atol=1e-5)
-
-    def test_torch_compile(self):
-        if not self._test_torch_compile:
-            return
-
-        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
-        compiled_model = torch.compile(model, mode="max-autotune", fullgraph=True, dynamic=False)
-
-        model.to(torch_device)
-        with torch.no_grad():
-            model_output = model(**self.get_dummy_inputs()).sample
-
-        compiled_model.to(torch_device)
-        with torch.no_grad():
-            compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample
-
-        model_output = model_output.detach().float().cpu().numpy()
-        compiled_model_output = compiled_model_output.detach().float().cpu().numpy()
-
-        max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten())
-        assert max_diff < 1e-3
-
-    def test_device_map_error(self):
-        with self.assertRaises(ValueError):
-            _ = self.model_cls.from_pretrained(
-                **self.get_dummy_model_init_kwargs(), device_map={0: "8GB", "cpu": "16GB"}
-            )
-
-
-class FluxTransformerQuantoMixin(QuantoBaseTesterMixin):
-    model_id = "hf-internal-testing/tiny-flux-transformer"
-    model_cls = FluxTransformer2DModel
-    pipeline_cls = FluxPipeline
-    torch_dtype = torch.bfloat16
-    keep_in_fp32_module = "proj_out"
-    modules_to_not_convert = ["proj_out"]
-    _test_torch_compile = False
-
-    def get_dummy_inputs(self):
-        return {
-            "hidden_states": torch.randn((1, 4096, 64), generator=torch.Generator("cpu").manual_seed(0)).to(
-                torch_device, self.torch_dtype
-            ),
-            "encoder_hidden_states": torch.randn(
-                (1, 512, 4096),
-                generator=torch.Generator("cpu").manual_seed(0),
-            ).to(torch_device, self.torch_dtype),
-            "pooled_projections": torch.randn(
-                (1, 768),
-                generator=torch.Generator("cpu").manual_seed(0),
-            ).to(torch_device, self.torch_dtype),
-            "timestep": torch.tensor([1]).to(torch_device, self.torch_dtype),
-            "img_ids": torch.randn((4096, 3), generator=torch.Generator("cpu").manual_seed(0)).to(
-                torch_device, self.torch_dtype
-            ),
-            "txt_ids": torch.randn((512, 3), generator=torch.Generator("cpu").manual_seed(0)).to(
-                torch_device, self.torch_dtype
-            ),
-            "guidance": torch.tensor([3.5]).to(torch_device, self.torch_dtype),
-        }
-
-    def get_dummy_training_inputs(self, device=None, seed: int = 0):
-        batch_size = 1
-        num_latent_channels = 4
-        num_image_channels = 3
-        height = width = 4
-        sequence_length = 48
-        embedding_dim = 32
-
-        torch.manual_seed(seed)
-        hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(device, dtype=torch.bfloat16)
-
-        torch.manual_seed(seed)
-        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(
-            device, dtype=torch.bfloat16
-        )
-
-        torch.manual_seed(seed)
-        pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(device, dtype=torch.bfloat16)
-
-        torch.manual_seed(seed)
-        text_ids = torch.randn((sequence_length, num_image_channels)).to(device, dtype=torch.bfloat16)
-
-        torch.manual_seed(seed)
-        image_ids = torch.randn((height * width, num_image_channels)).to(device, dtype=torch.bfloat16)
-
-        timestep = torch.tensor([1.0]).to(device, dtype=torch.bfloat16).expand(batch_size)
-
-        return {
-            "hidden_states": hidden_states,
-            "encoder_hidden_states": encoder_hidden_states,
-            "pooled_projections": pooled_prompt_embeds,
-            "txt_ids": text_ids,
-            "img_ids": image_ids,
-            "timestep": timestep,
-        }
-
-    def test_model_cpu_offload(self):
-        init_kwargs = self.get_dummy_init_kwargs()
-        transformer = self.model_cls.from_pretrained(
-            "hf-internal-testing/tiny-flux-pipe",
-            quantization_config=QuantoConfig(**init_kwargs),
-            subfolder="transformer",
-            torch_dtype=torch.bfloat16,
-        )
-        pipe = self.pipeline_cls.from_pretrained(
-            "hf-internal-testing/tiny-flux-pipe", transformer=transformer, torch_dtype=torch.bfloat16
-        )
-        pipe.enable_model_cpu_offload(device=torch_device)
-        _ = pipe("a cat holding a sign that says hello", num_inference_steps=2)
-
-    def test_training(self):
-        quantization_config = QuantoConfig(**self.get_dummy_init_kwargs())
-        quantized_model = self.model_cls.from_pretrained(
-            "hf-internal-testing/tiny-flux-pipe",
-            subfolder="transformer",
-            quantization_config=quantization_config,
-            torch_dtype=torch.bfloat16,
-        ).to(torch_device)
-
-        for param in quantized_model.parameters():
-            # freeze the model as only adapter layers will be trained
-            param.requires_grad = False
-            if param.ndim == 1:
-                param.data = param.data.to(torch.float32)
-
-        for _, module in quantized_model.named_modules():
-            if isinstance(module, Attention):
-                module.to_q = LoRALayer(module.to_q, rank=4)
-                module.to_k = LoRALayer(module.to_k, rank=4)
-                module.to_v = LoRALayer(module.to_v, rank=4)
-
-        with torch.amp.autocast(str(torch_device), dtype=torch.bfloat16):
-            inputs = self.get_dummy_training_inputs(torch_device)
-            output = quantized_model(**inputs)[0]
-            output.norm().backward()
-
-        for module in quantized_model.modules():
-            if isinstance(module, LoRALayer):
-                self.assertTrue(module.adapter[1].weight.grad is not None)
-
-
-class FluxTransformerFloat8WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
-    expected_memory_reduction = 0.3
-
-    def get_dummy_init_kwargs(self):
-        return {"weights_dtype": "float8"}
-
-
-class FluxTransformerInt8WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
-    expected_memory_reduction = 0.3
-    _test_torch_compile = True
-
-    def get_dummy_init_kwargs(self):
-        return {"weights_dtype": "int8"}
-
-
-class FluxTransformerInt4WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
-    expected_memory_reduction = 0.55
-
-    def get_dummy_init_kwargs(self):
-        return {"weights_dtype": "int4"}
-
-
-class FluxTransformerInt2WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
-    expected_memory_reduction = 0.65
-
-    def get_dummy_init_kwargs(self):
-        return {"weights_dtype": "int2"}