Use CPU-friendlier FLUX GGUF profile

app.py

@@ -8,12 +8,10 @@ from PIL import Image, ImageDraw, ImageFont, ImageOps
 
 
 MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"
-GGUF_URL = (
-    "https://huggingface.co/unsloth/FLUX.2-klein-4B-GGUF/resolve/main/"
-    "flux-2-klein-4b-Q4_K_M.gguf"
-)
+GGUF_BASE_URL = "https://huggingface.co/unsloth/FLUX.2-klein-4B-GGUF/resolve/main"
 MAX_SEED = 2_147_483_647
-MAX_GENERATION_EDGE = 1024
+GPU_MAX_GENERATION_EDGE = 1024
+CPU_MAX_GENERATION_EDGE = 512
 MIN_GENERATION_EDGE = 256
 SIZE_STEP = 32
 
@@ -50,6 +48,15 @@ def _dtype() -> torch.dtype:
     return torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
 
+def _gguf_url() -> str:
+    filename = "flux-2-klein-4b-Q4_K_M.gguf" if torch.cuda.is_available() else "flux-2-klein-4b-Q2_K.gguf"
+    return f"{GGUF_BASE_URL}/{filename}"
+
+
+def _max_generation_edge() -> int:
+    return GPU_MAX_GENERATION_EDGE if torch.cuda.is_available() else CPU_MAX_GENERATION_EDGE
+
+
 def get_pipeline() -> Flux2KleinPipeline:
     global PIPELINE
 
@@ -62,7 +69,7 @@ def get_pipeline() -> Flux2KleinPipeline:
 
         quantization_config = GGUFQuantizationConfig(compute_dtype=_dtype())
         transformer = Flux2Transformer2DModel.from_single_file(
-            GGUF_URL,
+            _gguf_url(),
             config=MODEL_ID,
             subfolder="transformer",
             quantization_config=quantization_config,
@@ -75,10 +82,12 @@ def get_pipeline() -> Flux2KleinPipeline:
             torch_dtype=_dtype(),
         )
         pipe.vae.enable_slicing()
+        pipe.vae.enable_tiling()
 
         if torch.cuda.is_available():
             pipe.enable_model_cpu_offload()
         else:
+            pipe.enable_attention_slicing()
             pipe.to("cpu")
 
         pipe.set_progress_bar_config(disable=True)
@@ -93,7 +102,8 @@ def _round_to_step(value: int, step: int = SIZE_STEP) -> int:
 def _generation_size(image: Image.Image) -> tuple[int, int]:
     width, height = image.size
     longest_edge = max(width, height)
-    scale = min(1.0, MAX_GENERATION_EDGE / longest_edge) if longest_edge else 1.0
+    max_generation_edge = _max_generation_edge()
+    scale = min(1.0, max_generation_edge / longest_edge) if longest_edge else 1.0
 
     resized_width = max(MIN_GENERATION_EDGE, int(width * scale))
     resized_height = max(MIN_GENERATION_EDGE, int(height * scale))
@@ -101,8 +111,8 @@ def _generation_size(image: Image.Image) -> tuple[int, int]:
     gen_width = _round_to_step(resized_width)
     gen_height = _round_to_step(resized_height)
 
-    gen_width = max(MIN_GENERATION_EDGE, min(MAX_GENERATION_EDGE, gen_width))
-    gen_height = max(MIN_GENERATION_EDGE, min(MAX_GENERATION_EDGE, gen_height))
+    gen_width = max(MIN_GENERATION_EDGE, min(max_generation_edge, gen_width))
+    gen_height = max(MIN_GENERATION_EDGE, min(max_generation_edge, gen_height))
     return gen_width, gen_height
 
 
@@ -212,6 +222,9 @@ with gr.Blocks(css=CSS) as demo:
                 Upload a normal photo and get a side-by-side comparison:
                 the left panel stays untouched, the right panel is regenerated
                 in a BSOD, computers, robots, and industrial sci-fi style.
+
+                On free CPU hardware, generation uses a lighter quant and smaller
+                working size, so higher step counts can be slow.
                 """
             )
 
@@ -251,7 +264,7 @@ with gr.Blocks(css=CSS) as demo:
                 minimum=1,
                 maximum=50,
                 step=1,
-                value=12,
+                value=4,
             )
             guidance_scale = gr.Slider(
                 label="Guidance scale",