Spaces:
Sleeping
Sleeping
| """ | |
| generation-space — SDXL + ControlNet-depth + IP-Adapter room renderer. | |
| /generate endpoint accepts: | |
| prompt : str | |
| negative_prompt : str | |
| depth_image_base64 : str (grayscale PNG, base64) | |
| ip_adapter_images_json : str (JSON array of product image URLs) | |
| controlnet_conditioning_scale : float | |
| ip_adapter_scale : float | |
| seed : int | |
| Returns: {"image_base64": str} (rendered room, JPEG, base64) | |
| ZeroGPU pattern: load all models on CPU at startup (no device_map, no cpu_offload), | |
| move to CUDA only inside @spaces.GPU. After inference, move back to CPU to free VRAM. | |
| VRAM budget on A10G (24 GB): | |
| SDXL base ~7 GB fp16 | |
| ControlNet-depth ~1.5 GB fp16 | |
| IP-Adapter ~0.5 GB (UNet patch) | |
| Total ~9 GB — comfortable on A10G | |
| """ | |
| import base64 | |
| import io | |
| import json | |
| import gradio as gr | |
| import httpx | |
| import numpy as np | |
| import spaces | |
| import torch | |
| from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline | |
| from PIL import Image | |
| SDXL_MODEL = "stabilityai/stable-diffusion-xl-base-1.0" | |
| CTRL_MODEL = "diffusers/controlnet-depth-sdxl-1.0" | |
| IPA_REPO = "h94/IP-Adapter" | |
| IPA_SUBDIR = "sdxl_models" | |
| IPA_WEIGHTS = "ip-adapter_sdxl.bin" | |
| # --------------------------------------------------------------------------- | |
| # Load on CPU at startup — ZeroGPU moves to GPU inside @spaces.GPU | |
| # --------------------------------------------------------------------------- | |
| print(f"Loading ControlNet: {CTRL_MODEL}...") | |
| controlnet = ControlNetModel.from_pretrained(CTRL_MODEL, torch_dtype=torch.float16) | |
| print(f"Loading SDXL: {SDXL_MODEL}...") | |
| pipe = StableDiffusionXLControlNetPipeline.from_pretrained( | |
| SDXL_MODEL, | |
| controlnet=controlnet, | |
| torch_dtype=torch.float16, | |
| use_safetensors=True, | |
| ) | |
| print(f"Loading IP-Adapter...") | |
| pipe.load_ip_adapter(IPA_REPO, subfolder=IPA_SUBDIR, weight_name=IPA_WEIGHTS) | |
| print("Generation pipeline ready.") | |
| # --------------------------------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------------------------------- | |
| def _decode_b64_image(b64: str) -> Image.Image: | |
| return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB") | |
| def _fetch_image(url: str) -> Image.Image: | |
| resp = httpx.get(url, timeout=15, follow_redirects=True) | |
| resp.raise_for_status() | |
| return Image.open(io.BytesIO(resp.content)).convert("RGB") | |
| OUTPUT_W = 1216 # standard SDXL landscape — 3:2, ~1MP, consistent across all inputs | |
| OUTPUT_H = 832 | |
| def _image_to_b64_jpeg(image: Image.Image) -> str: | |
| buf = io.BytesIO() | |
| image.save(buf, format="JPEG", quality=90) | |
| return base64.b64encode(buf.getvalue()).decode() | |
| # --------------------------------------------------------------------------- | |
| # Generation — GPU allocated for this function only | |
| # --------------------------------------------------------------------------- | |
| def generate( | |
| prompt: str, | |
| negative_prompt: str, | |
| depth_image_base64: str, | |
| ip_adapter_images_json: str, | |
| controlnet_conditioning_scale: float, | |
| ip_adapter_scale: float, | |
| seed: int, | |
| ) -> dict: | |
| import traceback | |
| try: | |
| device = "cuda" | |
| pipe.to(device) | |
| print(f"generate: prompt={prompt[:60]}") | |
| control_image = _decode_b64_image(depth_image_base64).resize((OUTPUT_W, OUTPUT_H)) | |
| print(f"generate: control_image resized to {OUTPUT_W}x{OUTPUT_H}") | |
| image_urls: list[str] = json.loads(ip_adapter_images_json) | |
| print(f"generate: fetching {len(image_urls)} IP-Adapter images") | |
| ip_images: list[Image.Image] = [] | |
| for url in image_urls: | |
| try: | |
| ip_images.append(_fetch_image(url).resize((224, 224))) | |
| print(f" fetched: {url[:60]}") | |
| except Exception as e: | |
| print(f" WARNING: could not fetch {url[:60]}: {e}") | |
| if not ip_images: | |
| print(" WARNING: no IP images fetched, using white fallback") | |
| ip_images = [Image.new("RGB", (224, 224), (255, 255, 255))] | |
| print(f"generate: {len(ip_images)} IP images ready, running SDXL...") | |
| pipe.set_ip_adapter_scale(float(ip_adapter_scale)) | |
| generator = torch.Generator(device=device).manual_seed(int(seed)) | |
| # width/height inferred from control_image size (OUTPUT_W x OUTPUT_H) | |
| result = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| image=control_image, | |
| ip_adapter_image=[ip_images], # wrap in list: multiple refs for 1 adapter | |
| num_inference_steps=30, | |
| guidance_scale=7.5, | |
| controlnet_conditioning_scale=float(controlnet_conditioning_scale), | |
| generator=generator, | |
| ).images[0] | |
| print("generate: SDXL done") | |
| pipe.to("cpu") | |
| torch.cuda.empty_cache() | |
| return {"image_base64": _image_to_b64_jpeg(result)} | |
| except Exception as e: | |
| traceback.print_exc() | |
| try: | |
| pipe.to("cpu") | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| raise ValueError(f"generate failed: {type(e).__name__}: {e}") from e | |
| # --------------------------------------------------------------------------- | |
| # Gradio interface | |
| # --------------------------------------------------------------------------- | |
| with gr.Blocks(title="Generation Space") as demo: | |
| gr.Markdown("## SDXL + ControlNet-Depth + IP-Adapter Room Generator") | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt_in = gr.Textbox(label="prompt", value="A Japandi living room, warm tones") | |
| neg_in = gr.Textbox(label="negative_prompt", value="ugly, blurry, unrealistic") | |
| depth_in = gr.Textbox(label="depth_image_base64", lines=3) | |
| ip_in = gr.Textbox(label="ip_adapter_images_json", value="[]") | |
| ctrl_scale_in = gr.Slider(0.0, 1.5, value=0.7, label="controlnet_conditioning_scale") | |
| ip_scale_in = gr.Slider(0.0, 1.0, value=0.5, label="ip_adapter_scale") | |
| seed_in = gr.Number(value=42, label="seed", precision=0) | |
| gr.Button("Generate").click( | |
| generate, | |
| inputs=[prompt_in, neg_in, depth_in, ip_in, ctrl_scale_in, ip_scale_in, seed_in], | |
| outputs=gr.JSON(label="result"), | |
| api_name="generate", | |
| ) | |
| demo.launch(show_error=True) | |