Switch to FLUX.2-klein-4B: ungated, bf16 (no bnb, no ZeroGPU patch), 4-step default

app.py

@@ -5,13 +5,12 @@ 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
 
 
-repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
+repo_id = "black-forest-labs/FLUX.2-klein-4B"
 
 # Take the pipeline apart into stages: each stage only loads the components it needs.
 blocks = ModularPipeline.from_pretrained(repo_id).blocks
@@ -43,8 +42,7 @@ def infer(
     randomize_seed=False,
     width=1024,
     height=1024,
-    guidance_scale=4.0,
-    num_inference_steps=28,
+    num_inference_steps=4,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -56,7 +54,6 @@ def infer(
     # `pipe.stream()` yields an event with the live pipeline state after every denoising step
     stream = pipe.stream(
         **text_embeddings,
-        guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         width=width,
         height=height,
@@ -93,7 +90,7 @@ css = """
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(
-            """# FLUX.2 [dev] — Live Preview with Modular Diffusers
+            """# FLUX.2 [klein] — Live Preview with Modular Diffusers
 Live latent preview powered by `pipe.stream()`: the pipeline yields its live state after every
 denoising step, and a preview pipeline built from flux2's own unpack + decode blocks renders it.
 No custom blocks, queues, or threads — see [huggingface/diffusers#14159](https://github.com/huggingface/diffusers/pull/14159).
@@ -141,29 +138,20 @@ No custom blocks, queues, or threads — see [huggingface/diffusers#14159](https
                     value=1024,
                 )
 
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance Scale",
-                    minimum=1,
-                    maximum=15,
-                    step=0.1,
-                    value=4.0,
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=28,
-                )
+            num_inference_steps = gr.Slider(
+                label="Number of inference steps",
+                minimum=1,
+                maximum=16,
+                step=1,
+                value=4,
+            )
 
         gr.Examples(examples=examples, fn=infer, inputs=[prompt], outputs=[result, seed], cache_examples=False)
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
         outputs=[result, seed],
     )
 

requirements.txt

@@ -2,5 +2,3 @@ accelerate
 git+https://github.com/huggingface/diffusers.git@refs/pull/14159/head
 torch
 transformers
-torchvision
-bitsandbytes

zerogpu_patch.py

@@ -1,29 +0,0 @@
-# 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