pin the quantisation stack that matches this NF4 checkpoint

app.py

@@ -1,93 +1,97 @@
-"""FLUX.2 [dev] 4-bit studio — text-to-image and multi-reference editing, self-contained.
-
-Runs the pre-quantized `diffusers/FLUX.2-dev-bnb-4bit` checkpoint: full dev quality at
-about a fifth of the footprint, with the text encoder loaded INSIDE this Space. That
-last part matters — the public dev demo delegates its text encoder to a separate Space,
-so it breaks for good whenever that other Space goes away.
-
-The /infer signature mirrors the common FLUX.2 Space API, so existing clients only need
-to change the space id.
-"""
-import os
-import random
-
-import gradio as gr
-import numpy as np
-import spaces
-import torch
-from diffusers import Flux2Pipeline
-from PIL import Image
-
-MODEL_ID = os.environ.get("FLUX2_MODEL", "diffusers/FLUX.2-dev-bnb-4bit")
-MAX_SEED = np.iinfo(np.int32).max
-MAX_SIDE = 1536
-
-pipe = Flux2Pipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
-pipe.enable_model_cpu_offload()          # 4-bit weights still want room to breathe
-
-
-def _duration(prompt, image_list, seed, width, height, num_inference_steps, guidance_scale):
-    """ZeroGPU reserves the slot before the call and passes it the same arguments, so
-    this signature must mirror _run's exactly."""
-    n = 1 + 0.6 * len(image_list or [])
-    return int(min(300, max(90, int(num_inference_steps) * 2.4 * n + 40)))
-
-
-@spaces.GPU(duration=_duration)
-def _run(prompt, image_list, seed, width, height, num_inference_steps, guidance_scale):
-    generator = torch.Generator(device="cuda").manual_seed(int(seed))
-    kwargs = dict(prompt=prompt, width=int(width), height=int(height),
-                  num_inference_steps=int(num_inference_steps),
-                  guidance_scale=float(guidance_scale), generator=generator)
-    if image_list:
-        kwargs["image"] = image_list
-    return pipe(**kwargs).images[0]
-
-
-def infer(prompt, input_images, seed=0, randomize_seed=False, width=1024, height=1024,
-          num_inference_steps=28, guidance_scale=4.0, prompt_upsampling=False,
-          progress=gr.Progress(track_tqdm=True)):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    images = []
-    for item in (input_images or []):
-        path = item[0] if isinstance(item, (list, tuple)) else item
-        if isinstance(path, dict):
-            path = path.get("image") or path.get("path") or path.get("name")
-        if path:
-            im = Image.open(path).convert("RGB")
-            im.thumbnail((MAX_SIDE, MAX_SIDE))
-            images.append(im)
-    width = min(int(width), MAX_SIDE)
-    height = min(int(height), MAX_SIDE)
-    out = _run(prompt, images, seed, width, height, num_inference_steps, guidance_scale)
-    return out, seed
-
-
-with gr.Blocks(title="FLUX.2 dev 4-bit studio") as demo:
-    gr.Markdown("## FLUX.2 [dev] 4-bit studio\nText-to-image and multi-reference editing. "
-                "Self-contained: the text encoder lives in this Space.")
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(label="Prompt", lines=4)
-            input_images = gr.Gallery(label="Reference images (optional)", type="filepath",
-                                      columns=4, height=200)
-            run = gr.Button("Generate", variant="primary")
-            with gr.Accordion("Settings", open=False):
-                seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
-                randomize_seed = gr.Checkbox(True, label="Randomize seed")
-                width = gr.Slider(256, MAX_SIDE, value=1024, step=32, label="Width")
-                height = gr.Slider(256, MAX_SIDE, value=1024, step=32, label="Height")
-                num_inference_steps = gr.Slider(4, 50, value=28, step=1, label="Steps")
-                guidance_scale = gr.Slider(1.0, 10.0, value=4.0, step=0.1, label="Guidance")
-                prompt_upsampling = gr.Checkbox(False, label="Prompt upsampling (unused)")
-        with gr.Column():
-            result = gr.Image(label="Result", type="pil")
-            used_seed = gr.Number(label="Seed used")
-
-    run.click(infer,
-              inputs=[prompt, input_images, seed, randomize_seed, width, height,
-                      num_inference_steps, guidance_scale, prompt_upsampling],
-              outputs=[result, used_seed], api_name="infer")
-
-demo.queue().launch(show_error=True)
+"""FLUX.2 [dev] 4-bit studio — text-to-image and multi-reference editing, self-contained.
+
+Runs the pre-quantized `diffusers/FLUX.2-dev-bnb-4bit` checkpoint: full dev quality at
+about a fifth of the footprint, with the text encoder loaded INSIDE this Space. That
+last part matters — the public dev demo delegates its text encoder to a separate Space,
+so it breaks for good whenever that other Space goes away.
+
+The /infer signature mirrors the common FLUX.2 Space API, so existing clients only need
+to change the space id.
+"""
+import os
+import random
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+from diffusers import Flux2Pipeline
+from PIL import Image
+
+MODEL_ID = os.environ.get("FLUX2_MODEL", "diffusers/FLUX.2-dev-bnb-4bit")
+MAX_SEED = np.iinfo(np.int32).max
+MAX_SIDE = 1536
+
+pipe = Flux2Pipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map=None,               # let the offload hook place things, not the loader
+)
+pipe.enable_model_cpu_offload()    # 4-bit weights still want room to breathe
+
+
+def _duration(prompt, image_list, seed, width, height, num_inference_steps, guidance_scale):
+    """ZeroGPU reserves the slot before the call and passes it the same arguments, so
+    this signature must mirror _run's exactly."""
+    n = 1 + 0.6 * len(image_list or [])
+    return int(min(300, max(90, int(num_inference_steps) * 2.4 * n + 40)))
+
+
+@spaces.GPU(duration=_duration)
+def _run(prompt, image_list, seed, width, height, num_inference_steps, guidance_scale):
+    generator = torch.Generator(device="cuda").manual_seed(int(seed))
+    kwargs = dict(prompt=prompt, width=int(width), height=int(height),
+                  num_inference_steps=int(num_inference_steps),
+                  guidance_scale=float(guidance_scale), generator=generator)
+    if image_list:
+        kwargs["image"] = image_list
+    return pipe(**kwargs).images[0]
+
+
+def infer(prompt, input_images, seed=0, randomize_seed=False, width=1024, height=1024,
+          num_inference_steps=28, guidance_scale=4.0, prompt_upsampling=False,
+          progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    images = []
+    for item in (input_images or []):
+        path = item[0] if isinstance(item, (list, tuple)) else item
+        if isinstance(path, dict):
+            path = path.get("image") or path.get("path") or path.get("name")
+        if path:
+            im = Image.open(path).convert("RGB")
+            im.thumbnail((MAX_SIDE, MAX_SIDE))
+            images.append(im)
+    width = min(int(width), MAX_SIDE)
+    height = min(int(height), MAX_SIDE)
+    out = _run(prompt, images, seed, width, height, num_inference_steps, guidance_scale)
+    return out, seed
+
+
+with gr.Blocks(title="FLUX.2 dev 4-bit studio") as demo:
+    gr.Markdown("## FLUX.2 [dev] 4-bit studio\nText-to-image and multi-reference editing. "
+                "Self-contained: the text encoder lives in this Space.")
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", lines=4)
+            input_images = gr.Gallery(label="Reference images (optional)", type="filepath",
+                                      columns=4, height=200)
+            run = gr.Button("Generate", variant="primary")
+            with gr.Accordion("Settings", open=False):
+                seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
+                randomize_seed = gr.Checkbox(True, label="Randomize seed")
+                width = gr.Slider(256, MAX_SIDE, value=1024, step=32, label="Width")
+                height = gr.Slider(256, MAX_SIDE, value=1024, step=32, label="Height")
+                num_inference_steps = gr.Slider(4, 50, value=28, step=1, label="Steps")
+                guidance_scale = gr.Slider(1.0, 10.0, value=4.0, step=0.1, label="Guidance")
+                prompt_upsampling = gr.Checkbox(False, label="Prompt upsampling (unused)")
+        with gr.Column():
+            result = gr.Image(label="Result", type="pil")
+            used_seed = gr.Number(label="Seed used")
+
+    run.click(infer,
+              inputs=[prompt, input_images, seed, randomize_seed, width, height,
+                      num_inference_steps, guidance_scale, prompt_upsampling],
+              outputs=[result, used_seed], api_name="infer")
+
+demo.queue().launch(show_error=True)