Spaces:
Running
on
Zero
Running
on
Zero
| import os, gc, random, re | |
| import gradio as gr | |
| import torch, spaces | |
| from PIL import Image, ImageFilter | |
| import numpy as np | |
| import qrcode | |
| from qrcode.constants import ERROR_CORRECT_H | |
| from diffusers import ( | |
| StableDiffusionPipeline, | |
| StableDiffusionControlNetPipeline, | |
| StableDiffusionControlNetImg2ImgPipeline, # NEW: img2img pipeline | |
| ControlNetModel, | |
| DPMSolverMultistepScheduler, | |
| ) | |
| # Optional: silence matplotlib cache warning in Spaces | |
| os.environ.setdefault("MPLCONFIGDIR", "/tmp/mpl") | |
| MODEL_ID = "runwayml/stable-diffusion-v1-5" | |
| CN_QRMON = "monster-labs/control_v1p_sd15_qrcode_monster" | |
| DTYPE = torch.float16 | |
| # ---------- helpers ---------- | |
| def snap8(x: int) -> int: | |
| x = max(256, min(1024, int(x))) | |
| return x - (x % 8) | |
| def normalize_color(c): | |
| if c is None: return "white" | |
| if isinstance(c, (tuple, list)): | |
| r, g, b = (int(max(0, min(255, round(float(x))))) for x in c[:3]); return (r, g, b) | |
| if isinstance(c, str): | |
| s = c.strip() | |
| if s.startswith("#"): return s | |
| m = re.match(r"rgba?\(\s*([0-9.]+)\s*,\s*([0-9.]+)\s*,\s*([0-9.]+)", s, re.IGNORECASE) | |
| if m: | |
| r = int(max(0, min(255, round(float(m.group(1)))))) | |
| g = int(max(0, min(255, round(float(m.group(2)))))) | |
| b = int(max(0, min(255, round(float(m.group(3)))))) | |
| return (r, g, b) | |
| return s | |
| return "white" | |
| def make_qr(url="http://www.mybirdfire.com", size=768, border=12, back_color="#808080", blur_radius=1.2): | |
| # Mid-gray background improves blending & scan rate with QR-Monster. | |
| qr = qrcode.QRCode(version=None, error_correction=ERROR_CORRECT_H, box_size=10, border=int(border)) | |
| qr.add_data(url.strip()); qr.make(fit=True) | |
| img = qr.make_image(fill_color="black", back_color=normalize_color(back_color)).convert("RGB") | |
| img = img.resize((int(size), int(size)), Image.NEAREST) | |
| if blur_radius and blur_radius > 0: | |
| img = img.filter(ImageFilter.GaussianBlur(radius=float(blur_radius))) | |
| return img | |
| def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: float = 0.6, feather: float = 1.0) -> Image.Image: | |
| """Gently push ControlNet-required blacks/whites for scannability.""" | |
| if strength <= 0: return stylized | |
| q = qr_img.convert("L") | |
| black_mask = q.point(lambda p: 255 if p < 128 else 0).filter(ImageFilter.GaussianBlur(radius=float(feather))) | |
| black = np.asarray(black_mask, dtype=np.float32) / 255.0 | |
| white = 1.0 - black | |
| s = np.asarray(stylized.convert("RGB"), dtype=np.float32) / 255.0 | |
| s = s * (1.0 - float(strength) * black[..., None]) # deepen blacks | |
| s = s + (1.0 - s) * (float(strength) * 0.85 * white[..., None]) # lift whites | |
| s = np.clip(s, 0.0, 1.0) | |
| return Image.fromarray((s * 255.0).astype(np.uint8), mode="RGB") | |
| # ---------- lazy pipelines (CPU-offloaded for ZeroGPU) ---------- | |
| _SD = None | |
| _CN_TXT2IMG = None | |
| _CN_IMG2IMG = None | |
| def _base_scheduler_for(pipe): | |
| pipe.scheduler = DPMSolverMultistepScheduler.from_config( | |
| pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++" | |
| ) | |
| pipe.enable_attention_slicing(); pipe.enable_vae_slicing(); pipe.enable_model_cpu_offload() | |
| return pipe | |
| def get_sd_pipe(): | |
| global _SD | |
| if _SD is None: | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype=DTYPE, | |
| safety_checker=None, | |
| use_safetensors=True, | |
| low_cpu_mem_usage=True, | |
| ) | |
| _SD = _base_scheduler_for(pipe) | |
| return _SD | |
| def get_qrmon_txt2img_pipe(): | |
| """(kept for completeness; not used in the two-stage flow)""" | |
| global _CN_TXT2IMG | |
| if _CN_TXT2IMG is None: | |
| cn = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True) | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| MODEL_ID, | |
| controlnet=cn, | |
| torch_dtype=DTYPE, | |
| safety_checker=None, | |
| use_safetensors=True, | |
| low_cpu_mem_usage=True, | |
| ) | |
| _CN_TXT2IMG = _base_scheduler_for(pipe) | |
| return _CN_TXT2IMG | |
| def get_qrmon_img2img_pipe(): | |
| """This is the pipeline we want for stage B.""" | |
| global _CN_IMG2IMG | |
| if _CN_IMG2IMG is None: | |
| cn = ControlNetModel.from_pretrained(CN_QRMON, torch_dtype=DTYPE, use_safetensors=True) | |
| pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( | |
| MODEL_ID, | |
| controlnet=cn, | |
| torch_dtype=DTYPE, | |
| safety_checker=None, | |
| use_safetensors=True, | |
| low_cpu_mem_usage=True, | |
| ) | |
| _CN_IMG2IMG = _base_scheduler_for(pipe) | |
| return _CN_IMG2IMG | |
| # ---------- ZeroGPU tasks ---------- | |
| def txt2img(prompt: str, negative: str, steps: int, cfg: float, width: int, height: int, seed: int): | |
| pipe = get_sd_pipe() | |
| w, h = snap8(width), snap8(height) | |
| if int(seed) < 0: | |
| seed = random.randint(0, 2**31 - 1) | |
| gen = torch.Generator(device="cuda").manual_seed(int(seed)) | |
| if torch.cuda.is_available(): torch.cuda.empty_cache() | |
| gc.collect() | |
| with torch.autocast(device_type="cuda", dtype=DTYPE): | |
| out = pipe( | |
| prompt=str(prompt), | |
| negative_prompt=str(negative or ""), | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(cfg), | |
| width=w, height=h, | |
| generator=gen, | |
| ) | |
| return out.images[0] | |
| def qr_stylize(url: str, style_prompt: str, negative: str, steps: int, cfg: float, | |
| size: int, border: int, back_color: str, blur: float, | |
| qr_weight: float, repair_strength: float, feather: float, seed: int, | |
| denoise: float = 0.45): | |
| s = snap8(size) | |
| # --- Stage A: base art (txt2img) --- | |
| sd = get_sd_pipe() | |
| if int(seed) < 0: | |
| seed = random.randint(0, 2**31 - 1) | |
| gen = torch.Generator(device="cuda").manual_seed(int(seed)) | |
| if torch.cuda.is_available(): torch.cuda.empty_cache() | |
| gc.collect() | |
| with torch.autocast(device_type="cuda", dtype=DTYPE): | |
| base = sd( | |
| prompt=str(style_prompt), # don't include "QR code" here | |
| negative_prompt=str(negative or ""), | |
| num_inference_steps=max(int(steps)//2, 12), | |
| guidance_scale=float(cfg), | |
| width=s, height=s, | |
| generator=gen, | |
| ).images[0] | |
| # Control image (QR) | |
| qr_img = make_qr(url=url, size=s, border=int(border), | |
| back_color=back_color, blur_radius=float(blur)) | |
| # --- Stage B: ControlNet img2img (QR Monster) --- | |
| pipe = get_qrmon_img2img_pipe() | |
| if torch.cuda.is_available(): torch.cuda.empty_cache() | |
| gc.collect() | |
| with torch.autocast(device_type="cuda", dtype=DTYPE): | |
| out = pipe( | |
| prompt=str(style_prompt), | |
| negative_prompt=str(negative or ""), | |
| image=base, # init image (img2img) | |
| control_image=qr_img, # control image (QR) | |
| strength=float(denoise), # 0.3–0.6 keeps composition | |
| controlnet_conditioning_scale=float(qr_weight), | |
| control_guidance_start=0.05, | |
| control_guidance_end=0.95, | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(cfg), | |
| width=s, height=s, | |
| generator=gen, | |
| ) | |
| img = out.images[0] | |
| img = enforce_qr_contrast(img, qr_img, strength=float(repair_strength), feather=float(feather)) | |
| return img, qr_img, base | |
| # ---------- UI ---------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# ZeroGPU Stable Diffusion + AI QR Codes (Monster v2)") | |
| with gr.Tab("Text → Image"): | |
| prompt = gr.Textbox(label="Prompt", value="a cozy reading nook, warm sunlight, cinematic lighting, highly detailed") | |
| negative = gr.Textbox(label="Negative (optional)", value="lowres, blurry, watermark, text") | |
| steps = gr.Slider(8, 40, value=28, step=1, label="Steps") | |
| cfg = gr.Slider(1.0, 12.0, value=7.0, step=0.5, label="CFG") | |
| width = gr.Slider(256, 1024, value=640, step=16, label="Width") | |
| height = gr.Slider(256, 1024, value=640, step=16, label="Height") | |
| seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)") | |
| out_img = gr.Image(label="Image", interactive=False) | |
| gr.Button("Generate").click(txt2img, [prompt, negative, steps, cfg, width, height, seed], out_img) | |
| with gr.Tab("QR Code Stylizer (ControlNet Monster — two-stage)"): | |
| url = gr.Textbox(label="URL/Text", value="http://www.mybirdfire.com") | |
| s_prompt = gr.Textbox(label="Style prompt (no 'QR code' needed)", value="baroque palace interior, intricate roots, dramatic lighting, ultra detailed") | |
| s_negative= gr.Textbox(label="Negative prompt", value="lowres, low contrast, blurry, jpeg artifacts, worst quality, watermark, text") | |
| size = gr.Slider(384, 1024, value=768, step=64, label="Canvas (px)") | |
| steps2 = gr.Slider(10, 60, value=28, step=1, label="Total steps") | |
| cfg2 = gr.Slider(1.0, 12.0, value=6.5, step=0.1, label="CFG") | |
| border = gr.Slider(4, 20, value=12, step=1, label="QR border (quiet zone)") | |
| back_col = gr.ColorPicker(value="#808080", label="QR background") | |
| blur = gr.Slider(0.0, 3.0, value=1.2, step=0.1, label="Soften control (blur)") | |
| qr_w = gr.Slider(0.6, 1.6, value=1.2, step=0.05, label="QR control weight") | |
| denoise = gr.Slider(0.2, 0.8, value=0.45, step=0.01, label="Denoising strength (Stage B)") | |
| repair = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Post repair strength") | |
| feather = gr.Slider(0.0, 3.0, value=1.0, step=0.1, label="Repair feather (px)") | |
| seed2 = gr.Number(value=-1, precision=0, label="Seed (-1 random)") | |
| final_img = gr.Image(label="Final stylized QR") | |
| ctrl_img = gr.Image(label="Control QR used") | |
| base_img = gr.Image(label="Base art (Stage A)") | |
| gr.Button("Stylize QR").click( | |
| qr_stylize, | |
| [url, s_prompt, s_negative, steps2, cfg2, size, border, back_col, blur, qr_w, repair, feather, seed2, denoise], | |
| [final_img, ctrl_img, base_img] | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=12).launch() | |