stable-diffusion-webui-forge/backend/operations.py

# Copyright (C) 2024 Forge - Establish the Structures
# Copyright (C) 2025 ComfyUI - where Optimization is Stolen
# Copyright (C) 2026 Haoming02 - Burnt the Kitchen

import contextlib
import time
from typing import Callable, Union

import torch

from backend import memory_management, stream, utils
from backend.args import args, dynamic_args
from backend.patcher.lora import merge_lora_to_weight


def scaled_dot_product_attention(q, k, v, *args, **kwargs):
    return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)


try:
    if torch.cuda.is_available() and memory_management.WINDOWS:
        import inspect

        from torch.nn.attention import SDPBackend, sdpa_kernel

        if "set_priority" in inspect.signature(sdpa_kernel).parameters:
            SDPA_BACKEND_PRIORITY = [
                SDPBackend.FLASH_ATTENTION,
                SDPBackend.EFFICIENT_ATTENTION,
                SDPBackend.MATH,
            ]

            SDPA_BACKEND_PRIORITY.insert(0, SDPBackend.CUDNN_ATTENTION)

            def scaled_dot_product_attention(q, k, v, *args, **kwargs):
                with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
                    return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)

except Exception:
    pass


# region Cast


def get_weight_and_bias(layer: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
    """Forge-Specific Function for on-the-fly LoRA"""
    loras: dict[str, list] = getattr(layer, "forge_online_loras", dict())

    weight: torch.Tensor = getattr(layer, "weight", None)
    weight_patches: list = loras.get("weight", None)
    if weight is not None and weight_patches is not None:
        weight = merge_lora_to_weight(patches=weight_patches, weight=weight, key="online_weight_lora", computation_dtype=weight.dtype)

    bias: torch.Tensor = getattr(layer, "bias", None)
    bias_patches: list = loras.get("bias", None)
    if bias is not None and bias_patches is not None:
        bias = merge_lora_to_weight(patches=bias_patches, weight=bias, key="online_bias_lora", computation_dtype=bias.dtype)

    return weight, bias


def weights_manual_cast(
    layer: Union[torch.nn.Module, "ForgeWeights"],
    x: torch.Tensor,
    *,
    dtype: torch.dtype = None,
    weight_fn: Callable = None,
    bias_fn: Callable = None,
    skip_weight_dtype: bool = False,
    skip_bias_dtype: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, tuple]:
    """
    Cast layer to input dtype/device
    * Reference: https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/ops.py#L210
    """

    target_dtype, target_device = x.dtype, x.device

    non_blocking = memory_management.device_supports_non_blocking(target_device)
    weight, bias = None, None

    weight_has_function: bool = len(layer.weight_function) > 0 or weight_fn is not None
    bias_has_function: bool = len(layer.bias_function) > 0 or bias_fn is not None

    weight_args = dict(device=target_device, dtype=dtype or target_dtype, non_blocking=non_blocking)
    if skip_weight_dtype or weight_has_function:
        weight_args.pop("dtype")

    bias_args = dict(device=target_device, dtype=target_dtype, non_blocking=non_blocking)
    if skip_bias_dtype or bias_has_function:
        bias_args.pop("dtype")

    if stream.should_use_stream():
        offload_stream = memory_management.get_offload_stream(target_device)
        context = stream.stream_context()(offload_stream)
    else:
        offload_stream = None
        context = None

    if layer.weight is not None:
        weight = memory_management.cast_to(
            layer.weight,
            **weight_args,
            copy=weight_has_function,
            context=context if layer.weight.device != target_device else None,
        )

    if layer.bias is not None:
        bias = memory_management.cast_to(
            layer.bias,
            **bias_args,
            copy=bias_has_function,
            context=context if layer.bias.device != target_device else None,
        )

    memory_management.sync_stream(target_device, offload_stream)

    weight_a = weight
    bias_a = bias

    if weight_has_function:
        if weight_fn is not None:
            weight = weight_fn(weight)
        if not skip_weight_dtype:
            weight = weight.to(dtype=target_dtype)
        for f in layer.weight_function:
            weight = f(weight)

    if bias_has_function:
        if bias_fn is not None:
            bias = bias_fn(bias)
        if not skip_bias_dtype:
            bias = bias.to(dtype=target_dtype)
        for f in layer.bias_function:
            bias = f(bias)

    loras: dict[str, list[torch.Tensor]] = getattr(layer, "forge_online_loras", dict())

    weight_patches = loras.get("weight", None)
    bias_patches = loras.get("bias", None)

    if weight is not None and weight_patches is not None:
        weight = merge_lora_to_weight(patches=weight_patches, weight=weight, key="online_weight_lora", computation_dtype=weight.dtype)

    if bias is not None and bias_patches is not None:
        bias = merge_lora_to_weight(patches=bias_patches, weight=bias, key="online_bias_lora", computation_dtype=bias.dtype)

    return weight, bias, (offload_stream, weight_a, bias_a)


@contextlib.contextmanager
def main_stream_worker(weight, bias, offload_stream: tuple[torch.Stream, torch.Tensor, torch.Tensor]):
    yield
    if offload_stream is None:
        return
    os, weight_a, bias_a = offload_stream
    if os is None:
        return
    if weight_a is not None:
        device = weight_a.device
    elif bias_a is not None:
        device = bias_a.device
    else:
        return
    os.wait_stream(memory_management.current_stream(device))


current_device: torch.device = None
current_dtype: torch.dtype = None
current_manual_cast_enabled: bool = False
current_bnb_dtype: str = None


# region Forge OPs


class ForgeWeights:
    parameters_manual_cast = False
    weight_function = []
    bias_function = []


class ForgeOperations:
    class Linear(torch.nn.Linear, ForgeWeights):
        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled

        def reset_parameters(self):
            return None

        def forward(self, x):
            if self.parameters_manual_cast:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return torch.nn.functional.linear(x, weight, bias)
            else:
                weight, bias = get_weight_and_bias(self)
                return torch.nn.functional.linear(x, weight, bias)

    class Conv1d(torch.nn.Conv1d, ForgeWeights):

        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled

        def reset_parameters(self):
            return None

        def forward(self, x):
            if self.parameters_manual_cast:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return self._conv_forward(x, weight, bias)
            else:
                weight, bias = get_weight_and_bias(self)
                return super()._conv_forward(x, weight, bias)

    class Conv2d(torch.nn.Conv2d, ForgeWeights):

        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled

        def reset_parameters(self):
            return None

        def forward(self, x):
            if self.parameters_manual_cast:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return self._conv_forward(x, weight, bias)
            else:
                weight, bias = get_weight_and_bias(self)
                return super()._conv_forward(x, weight, bias)

    class Conv3d(torch.nn.Conv3d, ForgeWeights):

        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled

        def reset_parameters(self):
            return None

        def _conv_forward(self, input, weight, bias, autopad=None, *args, **kwargs):
            if autopad == "causal_zero":
                weight = weight[:, :, -input.shape[2] :, :, :]
            if memory_management.NVIDIA_CONV3D_WORKAROUND and weight.dtype in (torch.float16, torch.bfloat16):
                out = torch.cudnn_convolution(input, weight, self.padding, self.stride, self.dilation, self.groups, benchmark=False, deterministic=False, allow_tf32=True)
                if bias is not None:
                    out += bias.reshape((1, -1) + (1,) * (out.ndim - 2))
                return out
            else:
                return super()._conv_forward(input, weight, bias, *args, **kwargs)

        def forward(self, x, *, autopad=None):
            if self.parameters_manual_cast or autopad is not None:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return self._conv_forward(x, weight, bias, autopad)
            else:
                weight, bias = get_weight_and_bias(self)
                return super()._conv_forward(x, weight, bias)

    class GroupNorm(torch.nn.GroupNorm, ForgeWeights):

        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled

        def reset_parameters(self):
            return None

        def forward(self, x):
            if self.parameters_manual_cast:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return torch.nn.functional.group_norm(x, self.num_groups, weight, bias, self.eps)
            else:
                return super().forward(x)

    class LayerNorm(torch.nn.LayerNorm, ForgeWeights):

        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled

        def reset_parameters(self):
            return None

        def forward(self, x):
            if self.parameters_manual_cast:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return torch.nn.functional.layer_norm(x, self.normalized_shape, weight, bias, self.eps)
            else:
                return super().forward(x)

    class RMSNorm(torch.nn.RMSNorm, ForgeWeights):

        def __init__(self, *args, add=False, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled
            self.bias = None
            self.add = add  # used by llama.py

        def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
            if prefix + "scale" in state_dict:  # Flux
                state_dict[prefix + "weight"] = state_dict.pop(prefix + "scale")
            super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)

        def reset_parameters(self):
            self.bias = None
            return None

        def forward(self, x):
            if self.parameters_manual_cast:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return torch.nn.functional.rms_norm(x, self.normalized_shape, (weight + 1.0) if self.add else weight, self.eps)
            elif self.add:
                return torch.nn.functional.rms_norm(x, self.normalized_shape, self.weight + 1.0, self.eps)
            else:
                return super().forward(x)

    class Embedding(torch.nn.Embedding, ForgeWeights):

        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            super().__init__(*args, **kwargs)
            self.parameters_manual_cast = current_manual_cast_enabled
            self.bias = None

        def reset_parameters(self):
            self.bias = None
            return None

        def forward(self, x):
            if self.parameters_manual_cast:
                weight, bias, signal = weights_manual_cast(self, x, skip_weight_dtype=True, skip_bias_dtype=True)
                with main_stream_worker(weight, bias, signal):
                    return torch.nn.functional.embedding(x, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse)
            else:
                return super().forward(x)


# region Int8

from backend.operations_int8 import (
    dequantize,
    int8_forward_dynamic,
    quantize_int8_tensorwise,
    stochastic_round_int8_delta,
)


class ForgeOperationsInt8(ForgeOperations):
    """Custom operations for INT8 tensorwise quantization"""

    excluded_names = []
    _is_prequantized = None

    class Linear(torch.nn.Linear, ForgeWeights):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.weight_scale = None
            self._is_quantized = False
            self.compute_dtype = torch.bfloat16
            self.lora_A = None
            self.lora_B = None
            self.lora_alpha = None

        def reset_parameters(self):
            return None

        def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
            weight_key = prefix + "weight"
            scale_key = prefix + "weight_scale"
            input_scale_key = prefix + "input_scale"
            bias_key = prefix + "bias"

            weight_scale = state_dict.pop(scale_key, None)
            state_dict.pop(prefix + "comfy_quant", None)
            weight_tensor = state_dict.pop(weight_key, None)

            # Pop input_scale to clean state_dict, but ignore it
            _ = state_dict.pop(input_scale_key, None)

            if weight_tensor is not None:
                if weight_tensor.dtype == torch.int8 and weight_scale is not None:
                    # Load Quantized
                    self._is_quantized = True
                    self.weight = torch.nn.Parameter(weight_tensor, requires_grad=False)
                    ForgeOperationsInt8._is_prequantized = True  # Found a quantized layer

                    if isinstance(weight_scale, torch.Tensor):
                        self.weight_scale = weight_scale.float().item() if weight_scale.numel() == 1 else weight_scale.float()
                    else:
                        self.weight_scale = float(weight_scale)

                elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32):
                    # Load High-Precision
                    # Detect if the model is pre-quantized if we don't know yet
                    if ForgeOperationsInt8._is_prequantized is None:
                        # Robust detection: scan keys and a sample of values
                        is_prequant = False
                        for k in state_dict.keys():
                            if "weight_scale" in k or "comfy_quant" in k:
                                is_prequant = True
                                break

                        if not is_prequant:
                            # Fallback: scan a sample of values for int8 tensors
                            for i, v in enumerate(state_dict.values()):
                                if i > 200:
                                    break  # Safety limit
                                if getattr(v, "dtype", None) == torch.int8:
                                    is_prequant = True
                                    break
                        ForgeOperationsInt8._is_prequantized = is_prequant

                    is_excluded = any(ex in prefix for ex in ForgeOperationsInt8.excluded_names)
                    is_dim1 = self.in_features == 1 or self.out_features == 1 or weight_tensor.ndim == 1

                    if is_excluded or is_dim1 or ForgeOperationsInt8._is_prequantized:
                        self._is_quantized = False
                        self.weight = torch.nn.Parameter(weight_tensor, requires_grad=False)
                    else:
                        # Quantize on the fly
                        device = torch.device("cuda") if torch.cuda.is_available() else weight_tensor.device
                        w_gpu = weight_tensor.to(device, non_blocking=True)
                        q_weight, q_scale = quantize_int8_tensorwise(w_gpu)

                        self.weight = torch.nn.Parameter(q_weight.cpu(), requires_grad=False)
                        self.weight_scale = q_scale.cpu() if isinstance(q_scale, torch.Tensor) else q_scale
                        self._is_quantized = True
                else:
                    self._is_quantized = False
                    self.weight = torch.nn.Parameter(weight_tensor, requires_grad=False)
            else:
                missing_keys.append(weight_key)

            bias_tensor = state_dict.pop(bias_key, None)
            if bias_tensor is not None:
                self.bias = torch.nn.Parameter(bias_tensor, requires_grad=False)
            else:
                self.bias = None

        def convert_weight(self, _weight, inplace=False):
            if not self._is_quantized:
                return _weight
            return self.weight

        def set_weight(self, out_weight, inplace_update=False, seed=0):
            if not self._is_quantized:
                if inplace_update:
                    self.weight.data.copy_(out_weight)
                else:
                    self.weight = torch.nn.Parameter(out_weight.to(self.weight.dtype), requires_grad=False)
                return

            if out_weight.dtype == torch.int8:
                if inplace_update:
                    self.weight.data.copy_(out_weight)
                else:
                    self.weight = torch.nn.Parameter(out_weight, requires_grad=False)
                return

            # Re-quantize if fallback occurred
            new_weight = stochastic_round_int8_delta(out_weight, self.weight_scale, seed)
            if inplace_update:
                self.weight.data.copy_(new_weight)
            else:
                self.weight = torch.nn.Parameter(new_weight, requires_grad=False)

        def set_bias(self, out_bias, inplace_update=False, seed=0):
            if out_bias is None:
                return
            if inplace_update:
                if self.bias is not None:
                    self.bias.data.copy_(out_bias)
            else:
                self.bias = torch.nn.Parameter(out_bias, requires_grad=False)

        def forward(self, x: torch.Tensor) -> torch.Tensor:
            """Fast forward using torch._int_mm for quantized weights."""

            if not self._is_quantized:
                weight, bias, signal = weights_manual_cast(self, x)
                with main_stream_worker(weight, bias, signal):
                    return torch.nn.functional.linear(x, weight, bias)

            # 1. Move weight/bias/scale to device (non_blocking)
            weight = self.weight.to(x.device, non_blocking=True)
            bias = self.bias.to(x.device, non_blocking=True) if self.bias is not None else None

            w_scale = self.weight_scale
            if isinstance(w_scale, torch.Tensor):
                w_scale = w_scale.to(x.device, non_blocking=True)

            compute_dtype = x.dtype if x.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16

            x_shape = x.shape
            x_2d = x.reshape(-1, x_shape[-1])

            if x_2d.shape[0] > 16:
                y = int8_forward_dynamic(x_2d, weight, w_scale, bias, compute_dtype)
            else:
                # Small batch fallback
                w_float = dequantize(weight, w_scale).to(x.dtype)
                y = torch.nn.functional.linear(x_2d, w_float, bias)

            # Dynamic LoRA Path
            if self.lora_A is not None and self.lora_B is not None:
                # Ensure LoRA tensors are on the same device as x
                lA = self.lora_A.to(x.device, non_blocking=True)
                lB = self.lora_B.to(x.device, non_blocking=True)

                lora_x = torch.nn.functional.linear(x_2d.to(lA.dtype), lA)
                lora_y = torch.nn.functional.linear(lora_x, lB)

                if self.lora_alpha is not None:
                    lora_y = lora_y * self.lora_alpha

                y = y + lora_y.to(y.dtype)

            return y.reshape(*x_shape[:-1], y.shape[-1])


# region BnB


if memory_management.bnb_enabled():

    from backend.operations_bnb import (
        ForgeLoader4Bit,
        functional_dequantize_4bit,
        functional_linear_4bits,
    )

    class ForgeOperationsBNB4bits(ForgeOperations):
        class Linear(ForgeLoader4Bit, ForgeWeights):
            def __init__(self, *args, **kwargs):
                super().__init__(device=current_device, dtype=current_dtype, quant_type=current_bnb_dtype)
                self.parameters_manual_cast = current_manual_cast_enabled

            def forward(self, x):
                if self.bias is not None and self.bias.dtype != x.dtype:
                    self.bias = utils.tensor2parameter(self.bias.to(x.dtype))

                if hasattr(self, "forge_online_loras"):
                    weight, bias, signal = weights_manual_cast(self, x, weight_fn=functional_dequantize_4bit, skip_bias_dtype=True)
                    with main_stream_worker(weight, bias, signal):
                        return torch.nn.functional.linear(x, weight, bias)

                if not self.parameters_manual_cast:
                    return functional_linear_4bits(x, self.weight, self.bias)
                elif not self.weight.bnb_quantized:
                    assert x.device.type == "cuda", "BnB must use CUDA as Computation Device"
                    layer_original_device = self.weight.device
                    self.weight = self.weight._quantize(x.device)
                    bias = self.bias.to(x.device) if self.bias is not None else None
                    out = functional_linear_4bits(x, self.weight, bias)
                    self.weight = self.weight.to(layer_original_device)
                    return out
                else:
                    weight, bias, signal = weights_manual_cast(self, x, skip_weight_dtype=True, skip_bias_dtype=True)
                    with main_stream_worker(weight, bias, signal):
                        return functional_linear_4bits(x, weight, bias)


# region GGUF


from backend.operations_gguf import dequantize_tensor


class ForgeOperationsGGUF(ForgeOperations):
    class Linear(torch.nn.Module, ForgeWeights):
        def __init__(self, *args, **kwargs):
            super().__init__()
            self.dummy = {"device": current_device, "dtype": current_dtype}
            self.weight = None
            self.bias = None

        def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
            if hasattr(self, "dummy"):
                if (computation_dtype := self.dummy["dtype"]) not in [torch.float16, torch.bfloat16]:
                    computation_dtype = torch.float16

                if prefix + "weight" in state_dict:
                    self.weight = state_dict[prefix + "weight"].to(device=self.dummy["device"])
                    self.weight.computation_dtype = computation_dtype
                if prefix + "bias" in state_dict:
                    self.bias = state_dict[prefix + "bias"].to(device=self.dummy["device"])
                    self.bias.computation_dtype = computation_dtype

                del self.dummy
            else:
                if prefix + "weight" in state_dict:
                    self.weight = state_dict[prefix + "weight"]
                if prefix + "bias" in state_dict:
                    self.bias = state_dict[prefix + "bias"]

        def _apply(self, fn, recurse=True):
            for k, p in self.named_parameters(recurse=False, remove_duplicate=True):
                setattr(self, k, utils.tensor2parameter(fn(p)))
            return self

        def forward(self, x):
            if self.bias is not None and self.bias.dtype != x.dtype:
                self.bias = utils.tensor2parameter(dequantize_tensor(self.bias).to(x.dtype))
            if self.weight is not None and self.weight.dtype != x.dtype and getattr(self.weight, "gguf_cls", None) is None:
                self.weight = utils.tensor2parameter(self.weight.to(x.dtype))

            weight, bias, signal = weights_manual_cast(self, x, weight_fn=dequantize_tensor, skip_bias_dtype=True)
            with main_stream_worker(weight, bias, signal):
                return torch.nn.functional.linear(x, weight, bias)

    class Conv2d(torch.nn.Conv2d, ForgeWeights):
        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.dummy = {"device": current_device, "dtype": current_dtype}
            self.weight = None
            self.bias = None

        def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
            if hasattr(self, "dummy"):
                if (computation_dtype := self.dummy["dtype"]) not in [torch.float16, torch.bfloat16]:
                    computation_dtype = torch.float16

                if prefix + "weight" in state_dict:
                    self.weight = state_dict[prefix + "weight"].to(device=self.dummy["device"])
                    self.weight.computation_dtype = computation_dtype
                if prefix + "bias" in state_dict:
                    self.bias = state_dict[prefix + "bias"].to(device=self.dummy["device"])
                    self.bias.computation_dtype = computation_dtype

                del self.dummy
            else:
                if prefix + "weight" in state_dict:
                    self.weight = state_dict[prefix + "weight"]
                if prefix + "bias" in state_dict:
                    self.bias = state_dict[prefix + "bias"]

        def _apply(self, fn, recurse=True):
            for k, p in self.named_parameters(recurse=False, remove_duplicate=True):
                setattr(self, k, utils.tensor2parameter(fn(p)))
            return self

        def forward(self, x):
            if self.bias is not None and self.bias.dtype != x.dtype:
                self.bias = utils.tensor2parameter(dequantize_tensor(self.bias).to(x.dtype))
            if self.weight is not None and self.weight.dtype != x.dtype and getattr(self.weight, "gguf_cls", None) is None:
                self.weight = utils.tensor2parameter(self.weight.to(x.dtype))

            weight, bias, signal = weights_manual_cast(self, x, weight_fn=dequantize_tensor, skip_bias_dtype=True)
            with main_stream_worker(weight, bias, signal):
                return super()._conv_forward(x, weight, bias)

    class Embedding(torch.nn.Embedding, ForgeWeights):
        def __init__(self, *args, **kwargs):
            kwargs["device"] = current_device
            kwargs["dtype"] = current_dtype
            super().__init__(*args, **kwargs)
            self.dummy = {"device": current_device, "dtype": current_dtype}
            self.weight = None
            self.bias = None

        def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
            if hasattr(self, "dummy"):
                if (computation_dtype := self.dummy["dtype"]) not in [torch.float16, torch.bfloat16]:
                    computation_dtype = torch.float16

                if prefix + "weight" in state_dict:
                    self.weight = state_dict[prefix + "weight"].to(device=self.dummy["device"])
                    self.weight.computation_dtype = computation_dtype

                del self.dummy
            else:
                if prefix + "weight" in state_dict:
                    self.weight = state_dict[prefix + "weight"]

        def _apply(self, fn, recurse=True):
            for k, p in self.named_parameters(recurse=False, remove_duplicate=True):
                setattr(self, k, utils.tensor2parameter(fn(p)))
            return self

        def reset_parameters(self):
            self.bias = None
            return None

        def forward(self, x):
            weight, bias, signal = weights_manual_cast(self, x, weight_fn=dequantize_tensor, skip_weight_dtype=True, skip_bias_dtype=True)
            with main_stream_worker(weight, bias, signal):
                return torch.nn.functional.embedding(x, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse)


# region fp8


from backend.operations_mixed_precision import (
    QuantizedTensor,
    TensorCoreFP8Layout,
    mixed_precision_ops,
)


def fp8_linear(self: torch.nn.Linear, input: torch.Tensor):
    # https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/ops.py#L615
    dtype = self.weight.dtype
    if dtype is not torch.float8_e4m3fn:
        return None

    input_dtype = input.dtype
    input_shape = input.shape
    tensor_3d = input.ndim == 3

    if tensor_3d:
        input = input.reshape(-1, input_shape[2])

    if input.ndim != 2:
        return None

    scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
    scale_input = torch.ones((), device=input.device, dtype=torch.float32)

    w, bias, signal = weights_manual_cast(self, input, dtype=dtype)

    with main_stream_worker(w, bias, signal):
        input = torch.clamp(input, min=-448, max=448, out=input)
        input_fp8 = input.to(dtype).contiguous()
        layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
        quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)

        layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
        quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
        o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)

    if tensor_3d:
        o = o.reshape((input_shape[0], input_shape[1], w.shape[0]))

    return o


class ForgeOperationsFP8(ForgeOperations):
    class Linear(ForgeOperations.Linear, ForgeWeights):
        def forward(self, x):
            try:
                if (out := fp8_linear(self, x)) is not None:
                    return out
            except Exception as e:
                memory_management.logger.error(f"Error during fp8_fast: {e}")

            return super().forward(x)


# region Tiled


class TiledOperations(ForgeOperations):
    class Conv2d(ForgeOperations.Conv2d):
        tile_size: int

        def __init__(self, *arg, **kwargs):
            super().__init__(*arg, **kwargs)
            self._3x1x1: bool = self.kernel_size == (3, 3) and self.stride == (1, 1) and self.padding == (1, 1)
            self.tile_size = args.tiled_conv2d

        @torch.inference_mode()
        def forward(self, x: torch.Tensor):
            if not self._3x1x1:
                return super().forward(x)

            B, C, H, W = x.shape

            if H <= self.tile_size and W <= self.tile_size:
                return super().forward(x)

            orig_forward = super().forward
            out_channels = self.out_channels if self.out_channels is not None else C

            out = torch.empty((B, out_channels, H, W), device=x.device, dtype=x.dtype, memory_format=torch.contiguous_format)
            non_blocking = memory_management.device_supports_non_blocking(x.device)

            for i in range(0, H, self.tile_size):
                i0 = max(i - 1, 0)
                i1 = min(i + self.tile_size + 1, H)
                pi = i - i0
                ph = min(self.tile_size, H - i)

                for j in range(0, W, self.tile_size):
                    j0 = max(j - 1, 0)
                    j1 = min(j + self.tile_size + 1, W)

                    tile = x[:, :, i0:i1, j0:j1]
                    tile_conv = orig_forward(tile)

                    pj = j - j0
                    pw = min(self.tile_size, W - j)

                    out[:, :, i : i + ph, j : j + pw].copy_(tile_conv[:, :, pi : pi + ph, pj : pj + pw], non_blocking=non_blocking)
                    del tile_conv

            return out


# region Pick OPs


@contextlib.contextmanager
def using_forge_operations(operations=None, device=None, dtype=None, manual_cast_enabled=False, bnb_dtype=None):
    global current_device, current_dtype, current_manual_cast_enabled, current_bnb_dtype

    current_device, current_dtype, current_manual_cast_enabled, current_bnb_dtype = device, dtype, manual_cast_enabled, bnb_dtype

    if isinstance(bnb_dtype, str):
        # https://github.com/BobJohnson24/ComfyUI-Flux2-INT8/blob/main/int8_unet_loader.py
        ForgeOperationsInt8._is_prequantized = None

        match bnb_dtype:
            case "Flux2K4B" | "Flux2K9B":
                ForgeOperationsInt8.excluded_names = ["img_in", "time_in", "guidance_in", "txt_in", "final_layer", "double_stream_modulation_img", "double_stream_modulation_txt", "single_stream_modulation"]
                operations = ForgeOperationsInt8
            case "ZImage":
                ForgeOperationsInt8.excluded_names = ["cap_embedder", "t_embedder", "x_embedder", "cap_pad_token", "context_refiner", "final_layer", "noise_refiner", "adaLN", "x_pad_token"]
                operations = ForgeOperationsInt8
            case "Chroma":
                ForgeOperationsInt8.excluded_names = ["distilled_guidance_layer", "final_layer", "img_in", "txt_in", "nerf_image_embedder", "nerf_blocks", "nerf_final_layer_conv", "__x0__", "nerf_final_layer_conv"]
                operations = ForgeOperationsInt8
            case "QwenImage":
                ForgeOperationsInt8.excluded_names = ["time_text_embed", "img_in", "norm_out", "proj_out", "txt_in"]
                operations = ForgeOperationsInt8
            case "WAN21_T2V" | "WAN21_I2V":
                ForgeOperationsInt8.excluded_names = ["patch_embedding", "text_embedding", "time_embedding", "time_projection" "head", "img_emb"]
                operations = ForgeOperationsInt8
    elif isinstance(bnb_dtype, dict):
        # https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/ops.py#L950

        _device = memory_management.get_torch_device()
        _dtype = torch.bfloat16 if memory_management.should_use_bf16(_device) else torch.float32
        fp8_compute = memory_management.supports_fp8_compute(_device)
        nvfp4_compute = memory_management.supports_nvfp4_compute(_device)
        mxfp8_compute = memory_management.supports_mxfp8_compute(_device)

        disabled = set()
        if not nvfp4_compute:
            disabled.add("nvfp4")
        if not mxfp8_compute:
            disabled.add("mxfp8")
        if not fp8_compute:
            disabled.add("float8_e4m3fn")
            disabled.add("float8_e5m2")

        _full: bool = bnb_dtype.pop("TE", False)  # https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/sd1_clip.py#L114
        operations = mixed_precision_ops(quant_config=bnb_dtype, compute_dtype=_dtype, full_precision_mm=_full, disabled=disabled)

    if operations is None:
        if bnb_dtype in ["gguf"]:
            operations = ForgeOperationsGGUF
        elif bnb_dtype in ["nf4", "fp4"]:
            assert memory_management.bnb_enabled(), 'Install the "bitsandbytes" package with --bnb'
            operations = ForgeOperationsBNB4bits
        elif bnb_dtype in ["vae"] and args.tiled_conv2d:
            memory_management.logger.info(f"Using TiledOperations ({args.tiled_conv2d}) for VAE")
            operations = TiledOperations
        elif dtype is torch.float8_e4m3fn and args.fast_fp8 and memory_management.supports_fp8_compute(memory_management.get_torch_device()):
            operations = ForgeOperationsFP8
        else:
            operations = ForgeOperations

    if operations is ForgeOperationsInt8:
        memory_management.logger.info("Quantizing to int8...")

    if dynamic_args.ops is None:
        dynamic_args.ops = str(operations.__name__)

    op_names = ("Linear", "Conv1d", "Conv2d", "Conv3d", "GroupNorm", "LayerNorm", "RMSNorm", "Embedding")
    backups = {op_name: getattr(torch.nn, op_name) for op_name in op_names}

    try:
        for op_name in op_names:
            setattr(torch.nn, op_name, getattr(operations, op_name))

        yield

    finally:
        for op_name in op_names:
            setattr(torch.nn, op_name, backups[op_name])


from functools import wraps


@contextlib.contextmanager
def automatic_memory_management():
    memory_management.free_memory(memory_required=3 * 1024 * 1024 * 1024, device=memory_management.get_torch_device())

    module_list: list[torch.nn.Module] = []

    original_init = torch.nn.Module.__init__
    original_to = torch.nn.Module.to

    @wraps(original_init)
    def patched_init(self, *args, **kwargs):
        module_list.append(self)
        return original_init(self, *args, **kwargs)

    @wraps(original_to)
    def patched_to(self, *args, **kwargs):
        module_list.append(self)
        return original_to(self, *args, **kwargs)

    try:
        torch.nn.Module.__init__ = patched_init
        torch.nn.Module.to = patched_to
        yield
    finally:
        torch.nn.Module.__init__ = original_init
        torch.nn.Module.to = original_to

    start = time.perf_counter()
    module_list = set(module_list)

    for module in module_list:
        module.cpu()

    memory_management.soft_empty_cache()
    end = time.perf_counter()

    memory_management.logger.debug(f"Automatic Memory Management: {len(module_list)} Modules in {(end - start):.2f} seconds")