Setup unsloth FLUX.2-klein-base-4B-GGUF local CPU Basic app

README.md

@@ -1,5 +1,5 @@
 ---
-title: Local Image Edit CPU Basic
+title: FLUX.2-klein-base-4B-GGUF CPU Basic
 colorFrom: gray
 colorTo: blue
 sdk: gradio
@@ -8,9 +8,11 @@ app_file: app.py
 pinned: false
 ---
 
-Space configurado para executar a edicao de imagem localmente no proprio CPU Basic.
+Space configurado para executar localmente no CPU Basic com:
+`unsloth/FLUX.2-klein-base-4B-GGUF`
 
 Notas:
 - Sem uso de Inference Providers remotos.
-- Primeira execucao pode demorar por download e load do modelo.
-- Se falhar por memoria/tempo, reduz o tamanho da imagem de entrada.
+- Primeira execucao demora (download + load do GGUF).
+- Default usa `flux-2-klein-base-4b-Q2_K.gguf` para reduzir RAM.
+- Se falhar, reduz imagem e steps.

app.py

@@ -1,19 +1,22 @@
 import os
+import random
 import threading
 
 import gradio as gr
 import torch
-from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
+from diffusers import Flux2KleinPipeline, Flux2Transformer2DModel, GGUFQuantizationConfig
+from huggingface_hub import hf_hub_download
 from PIL import Image
 
-MODEL_ID = os.getenv("MODEL_ID", "timbrooks/instruct-pix2pix")
+GGUF_REPO = "unsloth/FLUX.2-klein-base-4B-GGUF"
+BASE_REPO = "black-forest-labs/FLUX.2-klein-base-4B"
+GGUF_FILE = os.getenv("GGUF_FILE", "flux-2-klein-base-4b-Q2_K.gguf")
 MAX_SIDE = int(os.getenv("MAX_SIDE", "768"))
-NUM_STEPS = int(os.getenv("NUM_STEPS", "18"))
-GUIDANCE_SCALE = float(os.getenv("GUIDANCE_SCALE", "7.0"))
-IMAGE_GUIDANCE_SCALE = float(os.getenv("IMAGE_GUIDANCE_SCALE", "1.5"))
+DEFAULT_STEPS = int(os.getenv("DEFAULT_STEPS", "8"))
+DEFAULT_GUIDANCE = float(os.getenv("DEFAULT_GUIDANCE", "3.5"))
 
 _pipe = None
-_pipe_err = None
+_pipe_error = None
 _lock = threading.Lock()
 
 
@@ -23,72 +26,96 @@ def _prepare_image(img: Image.Image) -> Image.Image:
     m = max(w, h)
     if m > MAX_SIDE:
         s = MAX_SIDE / m
-        w, h = int(w * s), int(h * s)
-    w = max(64, (w // 8) * 8)
-    h = max(64, (h // 8) * 8)
+        w = int(w * s)
+        h = int(h * s)
+    w = max(256, (w // 32) * 32)
+    h = max(256, (h // 32) * 32)
     return img.resize((w, h), Image.Resampling.LANCZOS)
 
 
-def get_pipe() -> StableDiffusionInstructPix2PixPipeline:
-    global _pipe, _pipe_err
-    if _pipe_err:
-        raise RuntimeError(_pipe_err)
+def get_pipe() -> Flux2KleinPipeline:
+    global _pipe, _pipe_error
+    if _pipe_error is not None:
+        raise RuntimeError(_pipe_error)
     if _pipe is None:
         with _lock:
             if _pipe is None:
                 try:
-                    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
-                        MODEL_ID,
+                    gguf_path = hf_hub_download(repo_id=GGUF_REPO, filename=GGUF_FILE)
+                    qconfig = GGUFQuantizationConfig(compute_dtype=torch.float32)
+                    transformer = Flux2Transformer2DModel.from_single_file(
+                        gguf_path,
+                        quantization_config=qconfig,
                         torch_dtype=torch.float32,
-                        safety_checker=None,
                     )
-                    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+                    pipe = Flux2KleinPipeline.from_pretrained(
+                        BASE_REPO,
+                        transformer=transformer,
+                        torch_dtype=torch.float32,
+                    )
+                    pipe = pipe.to("cpu")
                     pipe.set_progress_bar_config(disable=True)
                     pipe.enable_attention_slicing()
-                    pipe = pipe.to("cpu")
+                    if hasattr(pipe, "enable_vae_slicing"):
+                        pipe.enable_vae_slicing()
                     _pipe = pipe
                 except Exception as e:
-                    _pipe_err = str(e)
+                    _pipe_error = str(e)
                     raise
     return _pipe
 
 
-def edit_image(image: Image.Image, prompt: str):
-    if image is None:
-        raise gr.Error("Carrega uma imagem primeiro.")
+def run_edit(image: Image.Image, prompt: str, steps: int, guidance: float, seed: int):
     if not prompt or not prompt.strip():
-        raise gr.Error("Escreve a instrução de edição.")
+        raise gr.Error("Escreve uma instrucao.")
 
     try:
         pipe = get_pipe()
     except Exception:
-        raise gr.Error("Falha ao carregar o modelo local no CPU Basic.")
+        raise gr.Error("Falha ao carregar FLUX.2-klein-base-4B-GGUF no CPU Basic.")
+
+    src = _prepare_image(image) if image is not None else None
+    if src is not None:
+        width, height = src.size
+    else:
+        width, height = 768, 768
+
+    if seed < 0:
+        seed = random.randint(0, 2**31 - 1)
+
+    generator = torch.Generator(device="cpu").manual_seed(seed)
 
-    src = _prepare_image(image)
     try:
         out = pipe(
             prompt=prompt.strip(),
             image=src,
-            num_inference_steps=NUM_STEPS,
-            guidance_scale=GUIDANCE_SCALE,
-            image_guidance_scale=IMAGE_GUIDANCE_SCALE,
+            height=height,
+            width=width,
+            num_inference_steps=max(1, int(steps)),
+            guidance_scale=float(guidance),
+            generator=generator,
         ).images[0]
         return out
     except Exception:
-        raise gr.Error("Falha durante a geração no CPU Basic. Tenta imagem menor.")
+        raise gr.Error("Falha na geracao. Tenta imagem menor e menos steps.")
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Image Edit local no CPU Basic")
-    gr.Markdown("Sem Inference Providers externos. Primeira execução pode demorar no download/load.")
+    gr.Markdown("# FLUX.2-klein-base-4B-GGUF local (CPU Basic)")
+    gr.Markdown("Modelo: unsloth/FLUX.2-klein-base-4B-GGUF (Q2_K por default).")
+
+    with gr.Row():
+        inp = gr.Image(type="pil", label="Imagem (opcional)")
+        out = gr.Image(type="pil", label="Resultado")
 
+    prompt = gr.Textbox(lines=3, label="Instrucao")
     with gr.Row():
-        input_img = gr.Image(type="pil", label="Imagem")
-        output_img = gr.Image(type="pil", label="Resultado")
+        steps = gr.Slider(minimum=1, maximum=20, value=DEFAULT_STEPS, step=1, label="Steps")
+        guidance = gr.Slider(minimum=1.0, maximum=8.0, value=DEFAULT_GUIDANCE, step=0.1, label="Guidance")
+        seed = gr.Number(value=-1, label="Seed (-1 aleatorio)")
 
-    prompt = gr.Textbox(lines=3, label="Instrução")
-    run = gr.Button("Editar")
-    run.click(edit_image, inputs=[input_img, prompt], outputs=output_img)
+    run_btn = gr.Button("Gerar")
+    run_btn.click(run_edit, inputs=[inp, prompt, steps, guidance, seed], outputs=out)
 
 
 if __name__ == "__main__":

requirements.txt

@@ -2,8 +2,9 @@
 
 gradio==5.29.1
 torch==2.5.1+cpu
-diffusers>=0.35.0
-transformers>=4.46.0,<5
+diffusers>=0.38.0
+transformers>=4.57.0,<5
 accelerate>=1.0.0
 safetensors>=0.4.5
+gguf>=0.13.0
 Pillow>=10.0.0