Spaces:
Paused
Paused
| """ | |
| 🎨 Klein LoRA Studio — restyle a photo into a named look, base vs your LoRA. | |
| Pick a named look (Scandinavian, Cottagecore, Watercolor…) and see a before→after | |
| restyle. Optionally load your own trained `.safetensors` LoRA so the look becomes | |
| a consistent *signature* that's truly yours. Image -> Image on FLUX.2 [klein] 4B. | |
| Build Small (Backyard AI). Built on the klein starter's ZeroGPU + LoRA pattern. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import random | |
| import time | |
| # --- ZeroGPU shim: import `spaces` BEFORE torch ----------------------------- | |
| try: | |
| import spaces # type: ignore | |
| GPU = spaces.GPU | |
| except Exception: # local / non-ZeroGPU fallback | |
| def GPU(*dargs, **dkwargs): # noqa: N802 | |
| if len(dargs) == 1 and callable(dargs[0]) and not dkwargs: | |
| return dargs[0] | |
| def wrap(fn): | |
| return fn | |
| return wrap | |
| import gradio as gr | |
| import torch | |
| from diffusers import Flux2KleinPipeline | |
| from PIL import Image | |
| MODEL_ID = "black-forest-labs/FLUX.2-klein-4B" # 4B, Apache 2.0, ungated | |
| STEPS = 4 | |
| GUIDANCE = 1.0 | |
| MAX_SEED = 2**31 - 1 | |
| pipe = None | |
| LOAD_ERR = "" | |
| try: | |
| print(f"Loading {MODEL_ID} on CPU…") | |
| pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16) | |
| print(" loaded.") | |
| except Exception as e: # noqa: BLE001 | |
| LOAD_ERR = str(e) | |
| print("Model load failed:", e) | |
| def klein_size(w: int, h: int, target_area: int = 1024 * 1024, divisor: int = 16): | |
| aspect = w / h | |
| nh = int((target_area / aspect) ** 0.5) | |
| nw = int(nh * aspect) | |
| nw = max(divisor, (nw // divisor) * divisor) | |
| nh = max(divisor, (nh // divisor) * divisor) | |
| return nw, nh | |
| LOOKS = { | |
| "Scandinavian": "Restyle as a clean Scandinavian interior, pale wood, white " | |
| "walls, muted neutral palette, soft daylight", | |
| "Cottagecore": "Restyle as a cozy cottagecore scene, warm florals, vintage " | |
| "textiles, soft golden light", | |
| "Mid-century modern": "Restyle as a mid-century modern interior, walnut tones, " | |
| "warm accent colours, retro furniture", | |
| "Cyberpunk neon": "Restyle with cyberpunk neon lighting, moody atmosphere, " | |
| "magenta and cyan glow, rain-slick reflections", | |
| "Watercolor": "Turn it into a soft watercolor painting, gentle washes, visible " | |
| "paper texture", | |
| "Risograph": "Restyle as a two-colour risograph print, grainy texture, bold " | |
| "flat inks", | |
| } | |
| def _apply_lora(p, path: str, scale: float) -> None: | |
| try: | |
| p.unload_lora_weights() | |
| except Exception: | |
| pass | |
| p.load_lora_weights(path, adapter_name="user") | |
| p.set_adapters(["user"], adapter_weights=[float(scale)]) | |
| def _unload(p) -> None: | |
| try: | |
| p.unload_lora_weights() | |
| except Exception: | |
| pass | |
| def restyle(input_image: Image.Image | None, look_key: str, lora_file, scale: float): | |
| if pipe is None: | |
| raise gr.Error(f"Model isn't loaded (this Space needs a GPU). {LOAD_ERR[:200]}") | |
| if input_image is None: | |
| raise gr.Error("Upload a photo first (or pick an example).") | |
| pipe.to("cuda") | |
| img = input_image.convert("RGB") | |
| w, h = klein_size(*img.size) | |
| if img.size != (w, h): | |
| img = img.resize((w, h), Image.LANCZOS) | |
| prompt = LOOKS.get(look_key, next(iter(LOOKS.values()))) | |
| seed = random.randint(0, MAX_SEED) | |
| t = time.time() | |
| try: | |
| _unload(pipe) | |
| if lora_file is not None: | |
| path = lora_file if isinstance(lora_file, str) else lora_file.name | |
| _apply_lora(pipe, path, scale) | |
| tag = f"+ LoRA {os.path.basename(path)} @ {scale}" | |
| else: | |
| tag = "base klein 4B" | |
| out = pipe( | |
| prompt=prompt, | |
| image=img, # keyword — `image` is positional-first | |
| width=w, | |
| height=h, | |
| num_inference_steps=STEPS, | |
| guidance_scale=GUIDANCE, | |
| generator=torch.Generator(device="cuda").manual_seed(seed), | |
| ).images[0] | |
| finally: | |
| _unload(pipe) # leave the shared pipe clean | |
| return img, out, f"**{look_key}** · {tag} · seed {seed} · {time.time() - t:.1f}s" | |
| THEME = gr.themes.Soft( | |
| font=["system-ui", "-apple-system", "Segoe UI", "Roboto", "Helvetica", "Arial", "sans-serif"], | |
| font_mono=["ui-monospace", "SFMono-Regular", "Consolas", "monospace"], | |
| ) | |
| CSS = """ | |
| footer {visibility: hidden;} | |
| .gradio-container, .gradio-container .prose, .gradio-container p, | |
| .gradio-container h1, .gradio-container h2, .gradio-container h3 { | |
| font-family: system-ui, -apple-system, "Segoe UI", Roboto, Helvetica, Arial, sans-serif !important; | |
| } | |
| """ | |
| with gr.Blocks(title="Klein LoRA Studio", theme=THEME, css=CSS) as demo: | |
| gr.Markdown( | |
| "# 🎨 Klein LoRA Studio\n" | |
| "Drop in a photo, pick a **named look**, and see a **before → after** " | |
| "restyle. Want the look to be *yours* and consistent across a whole set? " | |
| "Load a `.safetensors` LoRA you trained and it becomes your **signature** " | |
| "style. Powered by **FLUX.2 [klein] 4B** (4B params, Apache 2.0)." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| in_img = gr.Image(type="pil", label="Your photo", height=300) | |
| look = gr.Dropdown(list(LOOKS), value="Scandinavian", label="Look") | |
| lora = gr.File(label="Signature LoRA (optional .safetensors)", file_types=[".safetensors"]) | |
| scale = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="LoRA strength") | |
| btn = gr.Button("🎨 Restyle", variant="primary") | |
| with gr.Column(): | |
| with gr.Row(): | |
| before = gr.Image(label="Before", height=320) | |
| after = gr.Image(label="After", height=320) | |
| info = gr.Markdown() | |
| btn.click(restyle, [in_img, look, lora, scale], [before, after, info]) | |
| if __name__ == "__main__": | |
| # ssr_mode=False: Gradio-5 SSR commonly renders unstyled raw HTML on Spaces. | |
| demo.queue(max_size=8).launch( | |
| server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False | |
| ) | |