Work around bnb-4bit shape check on ZeroGPU's fake CUDA device

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import spaces
 import torch
 
+import zerogpu_patch  # noqa: F401  — bnb-4bit loading on ZeroGPU, see the module docstring
 from diffusers import ModularPipeline
 from diffusers.modular_pipelines import SequentialPipelineBlocks
 from diffusers.modular_pipelines.flux2.decoders import Flux2UnpackLatentsStep

zerogpu_patch.py

@@ -0,0 +1,29 @@
+# Workaround for loading a bnb-4bit diffusers model on ZeroGPU.
+#
+# On ZeroGPU the main process has no real GPU: the `spaces` package fakes the CUDA device, and its
+# proxy tensors expose `.shape` as a plain tuple. diffusers' bnb shape check calls
+# `current_param.shape.numel()`, which only exists on `torch.Size`, so loading the quantized
+# transformer crashes with `'tuple' object has no attribute 'numel'`
+# (src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py:206). Everything else in the bnb path works
+# under the ZeroGPU shim — the bnb-4bit text encoder loads fine through transformers, whose
+# equivalent check doesn't assume `torch.Size`.
+#
+# TODO: remove once the duck-typed check (`math.prod(shape)`) lands in diffusers.
+
+import math
+
+from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer
+
+
+def check_quantized_param_shape(self, param_name, current_param, loaded_param):
+    n = math.prod(current_param.shape)
+    inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
+    if tuple(loaded_param.shape) != inferred_shape:
+        raise ValueError(
+            f"Expected the flattened shape of the current param ({param_name}) to be {tuple(loaded_param.shape)} "
+            f"but is {inferred_shape}."
+        )
+    return True
+
+
+BnB4BitDiffusersQuantizer.check_quantized_param_shape = check_quantized_param_shape