Create app.py

app.py

@@ -0,0 +1,362 @@
+"""
+FLUX.2 Klein 9B (distilled) + optional uncensored TE + NSFW LoRAs
+Lighter than SDXL UnrealVision / Klein-base-50-steps.
+"""
+import os
+import gc
+import json
+import random
+import base64
+from io import BytesIO
+from pathlib import Path
+
+import gradio as gr
+from gradio import Server
+from fastapi.responses import HTMLResponse
+import numpy as np
+import spaces
+import torch
+from PIL import Image
+
+MAX_SEED = np.iinfo(np.int32).max
+LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("device:", device, "cuda:", torch.cuda.is_available())
+
+from diffusers import Flux2KleinPipeline
+
+dtype = torch.bfloat16
+
+# Distilled 9B = 4 steps, much lighter than base-9B (28–50 steps)
+MODEL_ID = os.getenv("KLEIN_MODEL", "black-forest-labs/FLUX.2-klein-9B")
+USE_UNCENSORED_TE = os.getenv("KLEIN_UNCENSORED_TE", "1").strip() not in ("0", "false", "no")
+UNCENSORED_TE_REPO = os.getenv(
+    "KLEIN_TE_REPO",
+    "ponpoke/flux2-klein-9b-uncensored-text-encoder",
+)
+
+ADAPTER_SPECS = {
+    "None": {
+        "repo": None,
+        "weights": None,
+        "adapter_name": None,
+        "default_strength": 0.0,
+        "hint": "No LoRA",
+    },
+    "NSFW-Solo": {
+        "repo": "diroverflo/FLux_Klein_9B_NSFW",
+        "weights": None,  # auto from repo
+        "adapter_name": "nsfw-solo",
+        "default_strength": 0.75,
+        "hint": "Solo NSFW focus (no faces in train set)",
+    },
+    "Consistency": {
+        "repo": "dx8152/Flux2-Klein-9B-Consistency",
+        "weights": None,
+        "adapter_name": "klein-consistency",
+        "default_strength": 0.8,
+        "hint": "Identity / edit consistency (not NSFW unlock)",
+    },
+}
+
+LOADED_ADAPTERS: set = set()
+ADAPTER_NAMES = list(ADAPTER_SPECS.keys())
+
+print(f"Loading Flux2KleinPipeline: {MODEL_ID}")
+pipe = Flux2KleinPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=dtype,
+)
+if torch.cuda.is_available():
+    try:
+        pipe.enable_model_cpu_offload()
+        print("cpu_offload OK")
+    except Exception as e:
+        print("offload fail:", e)
+        pipe = pipe.to(device)
+else:
+    pipe = pipe.to(device)
+
+# Best-effort uncensored text encoder swap
+if USE_UNCENSORED_TE:
+    try:
+        from transformers import AutoModel, AutoTokenizer
+
+        print(f"[TE] loading uncensored encoder: {UNCENSORED_TE_REPO}")
+        tok = AutoTokenizer.from_pretrained(
+            UNCENSORED_TE_REPO,
+            token=os.getenv("HF_TOKEN") or None,
+        )
+        te = AutoModel.from_pretrained(
+            UNCENSORED_TE_REPO,
+            torch_dtype=dtype,
+            token=os.getenv("HF_TOKEN") or None,
+        )
+        # Flux2Klein may expose text_encoder / tokenizer attributes
+        if hasattr(pipe, "text_encoder"):
+            pipe.text_encoder = te
+        if hasattr(pipe, "tokenizer"):
+            pipe.tokenizer = tok
+        print("[TE] uncensored text encoder attached (best-effort)")
+    except Exception as e:
+        print(f"[TE] skip uncensored TE ({type(e).__name__}: {e})")
+        print("[TE] continuing with official encoder + NSFW LoRAs")
+
+print("Pipeline ready.")
+
+
+def b64_to_pil_list(b64_json_str):
+    if not b64_json_str or str(b64_json_str).strip() in ("", "[]"):
+        return []
+    try:
+        b64_list = json.loads(b64_json_str)
+    except Exception:
+        return []
+    out = []
+    for b64_str in b64_list:
+        if not b64_str or not isinstance(b64_str, str):
+            continue
+        try:
+            if b64_str.startswith("data:image"):
+                _, data = b64_str.split(",", 1)
+            else:
+                data = b64_str
+            out.append(Image.open(BytesIO(base64.b64decode(data))).convert("RGB"))
+        except Exception as e:
+            print("decode error:", e)
+    return out
+
+
+def pil_to_b64_png(image: Image.Image) -> str:
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
+
+
+def update_dimensions(image: Image.Image, max_side: int = 1024):
+    if image is None:
+        return 1024, 1024
+    w, h = image.size
+    if w >= h:
+        nw = max_side
+        nh = int(nw * h / w)
+    else:
+        nh = max_side
+        nw = int(nh * w / h)
+    return max(16, (nw // 16) * 16), max(16, (nh // 16) * 16)
+
+
+app = Server(title="FLUX2-Klein-9B-NSFW")
+
+
+@app.api(name="check_safety", queue=False)
+def check_safety(prompt: str) -> dict:
+    return {"status": "ok"}
+
+
+@app.api(name="generate")
+@spaces.GPU(duration=120)
+def generate(
+    images_b64_json: str,
+    prompt: str,
+    lora_adapter: str,
+    seed: int,
+    randomize_seed: bool,
+    guidance_scale: float,
+    steps: int,
+    lora_strength: float = 0.75,
+) -> dict:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    if not prompt or not str(prompt).strip():
+        raise gr.Error("Prompt is empty.")
+
+    lora_adapter = (lora_adapter or "NSFW-Solo").strip()
+    spec = ADAPTER_SPECS.get(lora_adapter) or ADAPTER_SPECS["None"]
+    strength = float(lora_strength if lora_strength is not None else spec["default_strength"])
+    strength = max(0.0, min(strength, 1.5))
+
+    if spec["adapter_name"] is None or strength <= 0:
+        try:
+            pipe.disable_lora()
+        except Exception:
+            pass
+        print("--- LoRA off ---")
+    else:
+        name = spec["adapter_name"]
+        if name not in LOADED_ADAPTERS:
+            print(f"--- Loading LoRA {lora_adapter} from {spec['repo']} ---")
+            try:
+                kwargs = {"adapter_name": name}
+                if spec.get("weights"):
+                    kwargs["weight_name"] = spec["weights"]
+                pipe.load_lora_weights(spec["repo"], **kwargs)
+                LOADED_ADAPTERS.add(name)
+            except Exception as e:
+                raise gr.Error(f"LoRA load failed ({lora_adapter}): {e}")
+        try:
+            pipe.set_adapters([name], adapter_weights=[strength])
+        except Exception:
+            pipe.set_adapters([name])
+        print(f"--- LoRA {name} @ {strength} ---")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    generator = torch.Generator(
+        device="cuda" if torch.cuda.is_available() else "cpu"
+    ).manual_seed(seed)
+
+    # Distilled Klein defaults: 4 steps, guidance ~1.0–4.0
+    steps = int(max(1, min(int(steps or 4), 28)))
+    guidance_scale = float(max(0.0, min(float(guidance_scale or 1.0), 8.0)))
+
+    pil_images = b64_to_pil_list(images_b64_json)
+    if pil_images:
+        width, height = update_dimensions(pil_images[0])
+        processed = [im.resize((width, height), LANCZOS) for im in pil_images]
+        image_input = processed if len(processed) > 1 else processed[0]
+        print(f"I2I {width}x{height} steps={steps} cfg={guidance_scale}")
+    else:
+        width, height = 1024, 1024
+        image_input = None
+        print(f"T2I {width}x{height} steps={steps} cfg={guidance_scale}")
+
+    try:
+        kwargs = dict(
+            prompt=str(prompt).strip(),
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )
+        if image_input is not None:
+            kwargs["image"] = image_input
+        result = pipe(**kwargs).images[0]
+        return {
+            "image": pil_to_b64_png(result),
+            "seed": seed,
+            "status": "success",
+            "width": width,
+            "height": height,
+            "steps": steps,
+            "guidance_scale": guidance_scale,
+            "lora": lora_adapter,
+            "lora_strength": strength,
+        }
+    except Exception as e:
+        raise gr.Error(f"Inference failed: {type(e).__name__}: {e}")
+    finally:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+@app.get("/api/config")
+def client_config():
+    return {
+        "model": MODEL_ID,
+        "uncensored_te": USE_UNCENSORED_TE,
+        "loras": ADAPTER_NAMES,
+        "defaults": {
+            "steps": 4,
+            "guidance_scale": 1.0,
+            "lora": "NSFW-Solo",
+            "lora_strength": 0.75,
+        },
+        "notes": [
+            "Distilled Klein 9B — typically 4 steps",
+            "TE uncensored is best-effort; LoRA teaches NSFW concepts to DiT",
+            "Accept FLUX license on black-forest-labs/FLUX.2-klein-9B",
+            "HF_TOKEN if gated models need auth",
+        ],
+    }
+
+
+@app.get("/", response_class=HTMLResponse)
+async def homepage():
+    html_path = Path(__file__).resolve().parent / "index.html"
+    if html_path.exists():
+        return html_path.read_text(encoding="utf-8")
+    return """
+<!DOCTYPE html>
+<html lang="en"><head>
+<meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1"/>
+<title>FLUX.2 Klein 9B · NSFW</title>
+<style>
+:root{--bg:#0b0d12;--card:#141821;--fg:#e8eaed;--muted:#9aa0a6;--a:#a78bfa}
+*{box-sizing:border-box}body{margin:0;font-family:Inter,system-ui,sans-serif;background:var(--bg);color:var(--fg)}
+.wrap{max-width:920px;margin:0 auto;padding:24px 16px 48px}h1{margin:0 0 4px;font-size:1.45rem}
+.sub{color:var(--muted);margin-bottom:16px;font-size:.9rem}
+.card{background:var(--card);border:1px solid #252a36;border-radius:14px;padding:14px;margin-bottom:12px}
+label{display:block;font-size:.8rem;color:var(--muted);margin-bottom:6px}
+textarea,input,select{width:100%;background:#0a0c10;color:var(--fg);border:1px solid #2a2f3a;border-radius:10px;padding:10px}
+textarea{min-height:88px}.row{display:grid;grid-template-columns:1fr 1fr;gap:10px}
+button{width:100%;border:0;border-radius:12px;padding:12px;font-weight:700;cursor:pointer;
+background:linear-gradient(135deg,#7c3aed,#db2777);color:#fff;font-size:1rem}
+button:disabled{opacity:.5}#out{max-width:100%;border-radius:12px;margin-top:10px}
+#status{white-space:pre-wrap;font-size:.85rem;color:var(--muted)}.hint{font-size:.75rem;color:var(--muted);margin-top:6px}
+</style></head><body>
+<div class="wrap">
+  <h1>FLUX.2 Klein 9B</h1>
+  <div class="sub">Distilled · TE uncensored (best-effort) · NSFW LoRA</div>
+  <div class="card">
+    <label>Image (optional = edit / I2I)</label>
+    <input id="file" type="file" accept="image/*"/>
+  </div>
+  <div class="card">
+    <label>Prompt</label>
+    <textarea id="prompt">nude woman standing by a window, natural light, photorealistic skin, detailed</textarea>
+    <div class="hint">TE uncensored passa o prompt; LoRA ensina anatomia ao DiT.</div>
+  </div>
+  <div class="card row">
+    <div>
+      <label>LoRA</label>
+      <select id="lora"><option>NSFW-Solo</option><option>Consistency</option><option>None</option></select>
+    </div>
+    <div>
+      <label>LoRA strength</label>
+      <input id="str" type="number" value="0.75" min="0" max="1.5" step="0.05"/>
+    </div>
+    <div>
+      <label>Steps (distilled ≈ 4)</label>
+      <input id="steps" type="number" value="4" min="1" max="28"/>
+    </div>
+    <div>
+      <label>Guidance</label>
+      <input id="cfg" type="number" value="1.0" min="0" max="8" step="0.1"/>
+    </div>
+  </div>
+  <button id="run">Generate</button>
+  <div class="card"><div id="status">Ready.</div><img id="out"/></div>
+</div>
+<script>
+const $=id=>document.getElementById(id);
+function fileToB64(f){return new Promise((res,rej)=>{const r=new FileReader();r.onload=()=>res(r.result);r.onerror=rej;r.readAsDataURL(f);});}
+$("run").onclick=async()=>{
+  const btn=$("run");btn.disabled=true;$("status").textContent="Queuing…";$("out").removeAttribute("src");
+  try{
+    const images=[];const f=$("file").files[0];if(f)images.push(await fileToB64(f));
+    const body={data:[JSON.stringify(images),$("prompt").value,$("lora").value,0,true,
+      parseFloat($("cfg").value),parseInt($("steps").value,10),parseFloat($("str").value)]};
+    const res=await fetch("/gradio_api/call/generate",{method:"POST",headers:{"Content-Type":"application/json"},body:JSON.stringify(body)});
+    const ev=await res.json();if(!ev.event_id)throw new Error(JSON.stringify(ev));
+    const tr=await fetch(`/gradio_api/call/generate/${ev.event_id}`);
+    const text=await tr.text();let payload=null;
+    for(const line of text.trim().split("\\n")){if(line.startsWith("data:")){try{payload=JSON.parse(line.slice(5).trim());}catch{}}}
+    const out=Array.isArray(payload)?payload[0]:payload;
+    if(out&&out.image){$("out").src=out.image;$("status").textContent=`OK · seed=${out.seed} · steps=${out.steps} · lora=${out.lora}@${out.lora_strength}`;}
+    else $("status").textContent="Done: "+JSON.stringify(out).slice(0,400);
+  }catch(e){$("status").textContent="Error: "+e.message;}finally{btn.disabled=false;}
+};
+</script></body></html>
+"""
+
+
+if __name__ == "__main__":
+    print("Klein 9B distilled + NSFW LoRA — accept BFL license on the model page")
+    app.launch(show_error=True, mcp_server=True)