ronantakizawa/codereview-bench huggingface dataset:
Repository: huggingface/diffusers
File: src/diffusers/quantizers/gguf/utils.py (Python)

Code before:
            return type(result)(wrapped)
        else:
            return result


class GGUFLinear(nn.Linear):
    def __init__(
        self,
        in_features,
        out_features,
        bias=False,
        compute_dtype=None,
        device=None,
    ) -> None:
        super().__init__(in_features, out_features, bias, device)
        self.compute_dtype = compute_dtype
        self.device = device

    def forward(self, inputs: torch.Tensor):
        if ops is not None and self.weight.is_cuda and inputs.is_cuda:
            return self.forward_cuda(inputs)
        return self.forward_native(inputs)

    def forward_native(self, inputs: torch.Tensor):
        weight = dequantize_gguf_tensor(self.weight)
        weight = weight.to(self.compute_dtype)
        bias = self.bias.to(self.compute_dtype) if self.bias is not None else None

        output = torch.nn.functional.linear(inputs, weight, bias)
        return output

    def forward_cuda(self, inputs: torch.Tensor):
        quant_type = self.weight.quant_type
        output = _fused_mul_mat_gguf(inputs.to(self.compute_dtype), self.weight, quant_type)
        if self.bias is not None:
            output += self.bias.to(self.compute_dtype)
        return output


Review comment (suggestion):
This should be fairly safe as long as we get the same values (upto some tolerance) on both the native and kernels' variants.

Code after:
        elif type(result) in (list, tuple):
            # Preserve the original type (tuple or list)
            quant_type = cls._extract_quant_type(args)
            wrapped = [cls(x, quant_type=quant_type) if isinstance(x, torch.Tensor) else x for x in result]
            return type(result)(wrapped)
        else:
            return result


class GGUFLinear(nn.Linear):
    def __init__(
        self,
        in_features,
        out_features,
        bias=False,
        compute_dtype=None,
        device=None,
    ) -> None:
        super().__init__(in_features, out_features, bias, device)
        self.compute_dtype = compute_dtype
        self.device = device

    def forward(self, inputs: torch.Tensor):
        if ops is not None and self.weight.is_cuda and inputs.is_cuda:
            return self.forward_cuda(inputs)
        return self.forward_native(inputs)

    def forward_native(self, inputs: torch.Tensor):
        weight = dequantize_gguf_tensor(self.weight)
        weight = weight.to(self.compute_dtype)
        bias = self.bias.to(self.compute_dtype) if self.bias is not None else None

        output = torch.nn.functional.linear(inputs, weight, bias)
        return output

    def forward_cuda(self, inputs: torch.Tensor):
        quant_type = self.weight.quant_type
        output = _fused_mul_mat_gguf(inputs.to(self.compute_dtype), self.weight, quant_type)
        if self.bias is not None:
            output += self.bias.to(self.compute_dtype)
        return output
Repository: huggingface/diffusers
File: src/diffusers/models/attention_dispatch.py (Python)

Code before:
    is_causal: bool = False,
    window_size: Tuple[int, int] = (-1, -1),
    softcap: float = 0.0,
    deterministic: bool = False,
    return_attn_probs: bool = False,
) -> torch.Tensor:
    sig = inspect.signature(flash_attn_3_func)
    accepted = set(sig.parameters)
    params = {
        "q": query,
        "k": key,
        "v": value,
        "softmax_scale": scale,
        "causal": is_causal,
        "qv": None,
        "q_descale": None,
        "k_descale": None,
        "v_descale": None,
        "window_size": window_size,
        "attention_chunk": 0,
        "softcap": softcap,
        "num_splits": 1,
        "pack_gqa": None,
        "deterministic": deterministic,
        "sm_margin": 0,
    }
    kwargs = {}
    for name, value in params.items():
        if name not in accepted:
            logger.debug(f"{name} is not accepted by the `flash_attn_3_func` method, so it will be discarded.")
        else:
            kwargs[name] = value

    out, lse, *_ = flash_attn_3_func(**kwargs)
    return (out, lse) if return_attn_probs else out


@_AttentionBackendRegistry.register(
    AttentionBackendName._FLASH_VARLEN_3,
    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
)
def _flash_varlen_attention_3(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    cu_seqlens_q: Optional[torch.Tensor] = None,
    cu_seqlens_k: Optional[torch.Tensor] = None,
    max_seqlen_q: Optional[int] = None,
    max_seqlen_k: Optional[int] = None,
    scale: Optional[float] = None,
    is_causal: bool = False,

Review comment (suggestion):
`attention_chunk` isn't an acceptable argument in the VLLM implementation. We set it to `0`, anyway.

Code after:
        return_attn_probs=return_attn_probs,
    )
    out = out.unflatten(0, (batch_size, -1))

    return out


@_AttentionBackendRegistry.register(
    AttentionBackendName._FLASH_3,
    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
)
def _flash_attention_3(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: Optional[float] = None,
    is_causal: bool = False,
    window_size: Tuple[int, int] = (-1, -1),
    softcap: float = 0.0,
    deterministic: bool = False,
    return_attn_probs: bool = False,
) -> torch.Tensor:
    out, lse, *_ = flash_attn_3_func(
        q=query,
        k=key,
        v=value,
        softmax_scale=scale,
        causal=is_causal,
        qv=None,
        q_descale=None,
        k_descale=None,
        v_descale=None,
        window_size=window_size,
        attention_chunk=0,
        softcap=softcap,
        num_splits=1,
        pack_gqa=None,
        deterministic=deterministic,
        sm_margin=0,
    )
    return (out, lse) if return_attn_probs else out


@_AttentionBackendRegistry.register(
    AttentionBackendName._FLASH_3_HUB,
    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
)
def _flash_attention_3_hub(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
Repository: huggingface/diffusers
File: src/diffusers/utils/kernels_utils.py (Python)

Code before:
from .import_utils import is_kernels_available


_DEFAULT_HUB_ID_FA3 = "kernels-community/vllm-flash-attn3"


def _get_fa3_from_hub():
    if not is_kernels_available():
        return None
    else:
        from kernels import get_kernel

        try:
            flash_attn_3_hub = get_kernel(_DEFAULT_HUB_ID_FA3)
            return flash_attn_3_hub
        except Exception:
            return None


Review comment (bug):
Doing this kind of blocks the actual error from propagating to the user, so it will be hard to debug what the problem is. As this is the function trying to fetch from a remote server, handling any errors here makes more sense to me instead of returning a None. Maybe something like:

```python
def _get_fa3_from_hub():
    if not is_kernels_available():
        raise EnvironmentError("message about kernels package not being installed")
    else:
        from kernels import get_kernel

        try:
            flash_attn_3_hub = get_kernel(_DEFAULT_HUB_ID_FA3)
            return flash_attn_3_hub
        except Exception as e:
            logger.error(f"An error occurred while fetching kernel \"{_DEFAULT_HUB_ID_FA3}\" from the Hub: {e}")
            raise // re-raise so that user has full traceback 
```

IMO, this function should only be called if the user intends to use FA3 from the Hub, instead of being called always during the first import of attention dispatcher

Code after:
from ..utils import get_logger
from .import_utils import is_kernels_available


logger = get_logger(__name__)


_DEFAULT_HUB_ID_FA3 = "kernels-community/flash-attn3"


def _get_fa3_from_hub():
    if not is_kernels_available():
        return None
    else:
        from kernels import get_kernel

        try:
            # TODO: temporary revision for now. Remove when merged upstream into `main`.
            flash_attn_3_hub = get_kernel(_DEFAULT_HUB_ID_FA3, revision="fake-ops-return-probs")
            return flash_attn_3_hub
        except Exception as e:
            logger.error(f"An error occurred while fetching kernel '{_DEFAULT_HUB_ID_FA3}' from the Hub: {e}")
            raise
Repository: huggingface/diffusers
File: src/diffusers/models/attention_dispatch.py (Python)

Code before:


@_AttentionBackendRegistry.register(
    AttentionBackendName._FLASH_3_HUB,
    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
)
def _flash_attention_3_hub(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: Optional[float] = None,
    is_causal: bool = False,
    window_size: Tuple[int, int] = (-1, -1),
    softcap: float = 0.0,
    deterministic: bool = False,
    return_attn_probs: bool = False,
) -> torch.Tensor:
    out = flash_attn_3_func_hub(
        q=query,
        k=key,
        v=value,
        softmax_scale=scale,
        causal=is_causal,
        qv=None,
        q_descale=None,
        k_descale=None,
        v_descale=None,
        window_size=window_size,
        softcap=softcap,
        num_splits=1,
        pack_gqa=None,
        deterministic=deterministic,
        sm_margin=0,
    )
    lse = None
    return (out, lse) if return_attn_probs else out


@_AttentionBackendRegistry.register(
    AttentionBackendName._FLASH_VARLEN_3,
    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
)
def _flash_varlen_attention_3(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    cu_seqlens_q: Optional[torch.Tensor] = None,
    cu_seqlens_k: Optional[torch.Tensor] = None,
    max_seqlen_q: Optional[int] = None,
    max_seqlen_k: Optional[int] = None,
    scale: Optional[float] = None,

Review comment (suggestion):
Follow (internal) [this link](https://huggingface.slack.com/archives/C083BS3K2F9/p1756395601227149?thread_ts=1756282399.872609&cid=C083BS3K2F9) 

Code after:
        attention_chunk=0,
        softcap=softcap,
        num_splits=1,
        pack_gqa=None,
        deterministic=deterministic,
        sm_margin=0,
    )
    return (out, lse) if return_attn_probs else out


@_AttentionBackendRegistry.register(
    AttentionBackendName._FLASH_3_HUB,
    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
)
def _flash_attention_3_hub(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: Optional[float] = None,
    is_causal: bool = False,
    window_size: Tuple[int, int] = (-1, -1),
    softcap: float = 0.0,
    deterministic: bool = False,
    return_attn_probs: bool = False,
) -> torch.Tensor:
    out = flash_attn_3_func_hub(
        q=query,
        k=key,
        v=value,
        softmax_scale=scale,
        causal=is_causal,
        qv=None,
        q_descale=None,
        k_descale=None,
        v_descale=None,
        window_size=window_size,
        softcap=softcap,
        num_splits=1,
        pack_gqa=None,
        deterministic=deterministic,
        sm_margin=0,
        return_attn_probs=return_attn_probs,
    )
    # When `return_attn_probs` is True, the above returns a tuple of
    # actual outputs and lse.
    return (out[0], out[1]) if return_attn_probs else out


@_AttentionBackendRegistry.register(
    AttentionBackendName._FLASH_VARLEN_3,
    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
Repository: huggingface/diffusers
File: tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py (Python)

Code before:
            num_inference_steps=2,
            negative_prompt="",
        ).to_tuple()

        generator = torch.Generator(device="cpu").manual_seed(0)
        output = pipeline(
            image_embeds=image_emb,
            negative_image_embeds=zero_image_emb,
            hint=hint,
            generator=generator,
            num_inference_steps=2,
            output_type="np",
        )

        image = output.images[0]

        assert image.shape == (512, 512, 3)
        max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten())
        expected_slices = Expectations(
            {
                ("xpu", 3): 2e-3,
                ("cuda", 7): 2e-4,
            }
        )
        expected_max_diff = expected_slices.get_expectation()
        assert max_diff < expected_max_diff


Review comment (suggestion):

```suggestion
        expected_max_diffs = Expectations(
            {
                ("xpu", 3): 2e-3,
                ("cuda", 7): 2e-4,
            }
        )
        expected_max_diff = expected_max_diffs.get_expectation()
        assert max_diff < expected_max_diff
```

Code after:
            num_inference_steps=2,
            negative_prompt="",
        ).to_tuple()

        generator = torch.Generator(device="cpu").manual_seed(0)
        output = pipeline(
            image_embeds=image_emb,
            negative_image_embeds=zero_image_emb,
            hint=hint,
            generator=generator,
            num_inference_steps=2,
            output_type="np",
        )

        image = output.images[0]

        assert image.shape == (512, 512, 3)
        max_diff = numpy_cosine_similarity_distance(expected_image.flatten(), image.flatten())
        expected_max_diffs = Expectations(
            {
                ("xpu", 3): 2e-3,
                ("cuda", 7): 2e-4,
            }
        )
        expected_max_diff = expected_max_diffs.get_expectation()
        assert max_diff < expected_max_diff
Repository: huggingface/diffusers
File: tests/lora/test_lora_layers_z_image.py (Python)

Code before:
# coding=utf-8
# Copyright 2025 HuggingFace Inc.
#
# 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.
import os
import sys
import unittest

import numpy as np
import torch
from peft import LoraConfig
from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model

from diffusers import (
    AutoencoderKL,
    FlowMatchEulerDiscreteScheduler,
    ZImagePipeline,
    ZImageTransformer2DModel,
)

from ..testing_utils import floats_tensor, require_peft_backend, torch_device


# Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
torch.use_deterministic_algorithms(False)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
if hasattr(torch.backends, "cuda"):
    torch.backends.cuda.matmul.allow_tf32 = False


sys.path.append(".")

from .utils import PeftLoraLoaderMixinTests, check_if_lora_correctly_set  # noqa: E402

Review comment (suggestion):
We have to keep it like this:
https://github.com/huggingface/diffusers/blob/152f7ca357c066c4af3d1a58cdf17662ef5a2f87/tests/lora/test_lora_layers_flux.py#L47

Code after:
# coding=utf-8
# Copyright 2025 HuggingFace Inc.
#
# 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.
import sys
import unittest

import torch
from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model

from diffusers import (
    AutoencoderKL,
    FlowMatchEulerDiscreteScheduler,
    ZImagePipeline,
    ZImageTransformer2DModel,
)

from ..testing_utils import floats_tensor, is_peft_available, require_peft_backend


if is_peft_available():
    from peft import LoraConfig


sys.path.append(".")

from .utils import PeftLoraLoaderMixinTests  # noqa: E402


@unittest.skip(
    "ZImage LoRA tests are skipped due to non-deterministic behavior from complex64 RoPE operations "
    "and torch.empty padding tokens. LoRA functionality works correctly with real models."
)
@require_peft_backend
class ZImageLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
    pipeline_class = ZImagePipeline
Repository: huggingface/diffusers
File: src/diffusers/utils/logging.py (Python)

Code before:

        library_root_logger = _get_library_root_logger()
        library_root_logger.removeHandler(_default_handler)
        library_root_logger.setLevel(logging.NOTSET)
        _default_handler = None


def get_log_levels_dict() -> Dict[str, int]:
    return log_levels


def get_logger(name: Optional[str] = None) -> logging.Logger:
    """
    Return a logger with the specified name.

    This function is not supposed to be directly accessed unless you are writing a custom diffusers module.
    """

    if name is None:
        name = _get_library_name()

    _configure_library_root_logger()
    logger = logging.getLogger(name)
    _ensure_rank_zero_filter(logger)
    return logger


def get_verbosity() -> int:
    """
    Return the current level for the 🤗 Diffusers' root logger as an `int`.

    Returns:
        `int`:
            Logging level integers which can be one of:

            - `50`: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
            - `40`: `diffusers.logging.ERROR`
            - `30`: `diffusers.logging.WARNING` or `diffusers.logging.WARN`
            - `20`: `diffusers.logging.INFO`
            - `10`: `diffusers.logging.DEBUG`

    """

    _configure_library_root_logger()
    return _get_library_root_logger().getEffectiveLevel()


def set_verbosity(verbosity: int) -> None:
    """
    Set the verbosity level for the 🤗 Diffusers' root logger.


Review comment (suggestion):
Helps with `logging` from modules within `diffusers`, ensuring that they only log on rank-0 processes.

Code after:
            return

        library_root_logger = _get_library_root_logger()
        library_root_logger.removeHandler(_default_handler)
        library_root_logger.setLevel(logging.NOTSET)
        _default_handler = None


def get_log_levels_dict() -> Dict[str, int]:
    return log_levels


def get_logger(name: Optional[str] = None) -> logging.Logger:
    """
    Return a logger with the specified name.

    This function is not supposed to be directly accessed unless you are writing a custom diffusers module.
    """

    if name is None:
        name = _get_library_name()

    _configure_library_root_logger()
    logger = logging.getLogger(name)
    _ensure_rank_zero_filter(logger)
    return logger


def get_verbosity() -> int:
    """
    Return the current level for the 🤗 Diffusers' root logger as an `int`.

    Returns:
        `int`:
            Logging level integers which can be one of:

            - `50`: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
            - `40`: `diffusers.logging.ERROR`
            - `30`: `diffusers.logging.WARNING` or `diffusers.logging.WARN`
            - `20`: `diffusers.logging.INFO`
            - `10`: `diffusers.logging.DEBUG`

    """

    _configure_library_root_logger()
    return _get_library_root_logger().getEffectiveLevel()


def set_verbosity(verbosity: int) -> None:
    """
    Set the verbosity level for the 🤗 Diffusers' root logger.
Repository: huggingface/diffusers
File: src/diffusers/training_utils.py (Python)

Code before:
import contextlib
import copy
import gc
import math
import random
import re
import warnings
from contextlib import contextmanager
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Set, Type

import numpy as np
import torch
import torch.distributed as dist
from torch.distributed.fsdp import CPUOffload, ShardingStrategy
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial

from .models import UNet2DConditionModel
from .pipelines import DiffusionPipeline
from .schedulers import SchedulerMixin
from .utils import (
    convert_state_dict_to_diffusers,
    convert_state_dict_to_peft,
    deprecate,
    is_peft_available,
    is_torch_npu_available,
    is_torchvision_available,
    is_transformers_available,
)


if is_transformers_available():
    import transformers

    if transformers.integrations.deepspeed.is_deepspeed_zero3_enabled():
        import deepspeed

if is_peft_available():
    from peft import set_peft_model_state_dict


Review comment (suggestion):
We should guard this like so:

```
if getattr(torch, "distributed", None) is not None:
    import torch.distributed as dist
```

Same for FSDP.

Code after:
import contextlib
import copy
import gc
import math
import random
import re
import warnings
from contextlib import contextmanager
from functools import partial
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type, Union

import numpy as np
import torch


if getattr(torch, "distributed", None) is not None:
    from torch.distributed.fsdp import CPUOffload, ShardingStrategy
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

from .models import UNet2DConditionModel
from .pipelines import DiffusionPipeline
from .schedulers import SchedulerMixin
from .utils import (
    convert_state_dict_to_diffusers,
    convert_state_dict_to_peft,
    deprecate,
    is_accelerate_available,
    is_peft_available,
    is_torch_npu_available,
    is_torchvision_available,
    is_transformers_available,
)


if is_transformers_available():
    import transformers

    if transformers.integrations.deepspeed.is_deepspeed_zero3_enabled():
        import deepspeed