import gc import json import random import threading import traceback import gradio as gr import spaces import torch from PIL import Image, ImageDraw from diffusers import QwenImage21Pipeline BASE_MODEL = "Qwen/Qwen-Image-2.1" LORA_REPO = "retronic/Pixelship-2.0" ADAPTER_NAME = "pixelship" DTYPE = torch.bfloat16 pipe = None lock = threading.Lock() def load_pipeline(): global pipe if pipe is None: print("Loading Qwen Image 2.1...", flush=True) p = QwenImage21Pipeline.from_pretrained( BASE_MODEL, torch_dtype=DTYPE, low_cpu_mem_usage=True ) p.load_lora_weights(LORA_REPO, adapter_name=ADAPTER_NAME) p.set_adapters([ADAPTER_NAME], adapter_weights=[0.75]) p.enable_model_cpu_offload() pipe = p print("Pixelship pipeline ready", flush=True) return pipe def parse_plan(raw): data = json.loads(raw) if not isinstance(data, dict): raise ValueError("JSON plan must be an object") for key in ("prompt", "math", "objects", "background", "style"): if key not in data: raise ValueError(f"Missing JSON field: {key}") if not isinstance(data["objects"], list) or not isinstance(data["background"], dict): raise ValueError("objects must be a list and background must be an object") bg_objects = data["background"].get("objects", []) if not isinstance(bg_objects, list): raise ValueError("background.objects must be a list") for obj in data["objects"] + bg_objects: if not isinstance(obj, dict): raise ValueError("Each object must be a JSON object") for key in ("description", "x", "y", "width", "height"): if key not in obj: raise ValueError(f"Missing object field: {key}") for key in ("x", "y", "width", "height"): v = float(obj[key]) if not 0 <= v <= 1: raise ValueError(f"{key} must be between 0 and 1") if float(obj["x"]) + float(obj["width"]) > 1.001 or float(obj["y"]) + float(obj["height"]) > 1.001: raise ValueError("Object bounds must fit within the image") return data def draw_map(plan, width, height): w, h = int(width), int(height) im = Image.new("RGB", (w, h), "#20232a") d = ImageDraw.Draw(im) for x in range(0, w, max(1, w // 10)): d.line((x, 0, x, h), fill="#353a43") for y in range(0, h, max(1, h // 10)): d.line((0, y, w, y), fill="#353a43") entries = [(o, "#4ecdc4") for o in plan["background"].get("objects", [])] entries += [(o, "#ffce54") for o in plan["objects"]] for i, (obj, color) in enumerate(entries, 1): x0 = round(float(obj["x"]) * w) y0 = round(float(obj["y"]) * h) x1 = round((float(obj["x"]) + float(obj["width"])) * w) y1 = round((float(obj["y"]) + float(obj["height"])) * h) d.rectangle((x0, y0, x1, y1), outline=color, width=max(1, w // 500)) d.text((min(w - 1, x0 + 2), max(0, y0 - 13)), f'{i}: {str(obj["description"])[:38]}', fill=color, stroke_width=1, stroke_fill="#101010") return im @spaces.GPU(duration=120) def generate(prompt, input_image, negative_prompt, width, height, steps, true_cfg, lora_strength, seed, random_seed, use_json, plan_text, show_map, progress=gr.Progress()): if not str(prompt or "").strip() and not (use_json and str(plan_text or "").strip()): raise gr.Error("Enter a prompt or supply a JSON plan") with lock: try: plan = None actual_prompt = str(prompt or "") if use_json: if not str(plan_text or "").strip(): raise gr.Error("Paste a JSON plan into the JSON box") plan = parse_plan(plan_text) actual_prompt = json.dumps(plan, ensure_ascii=False) progress(0.05, desc="Loading Pixelship 2.0") p = load_pipeline() seed = random.randint(0, 2**32 - 1) if random_seed else int(seed) p.set_adapters([ADAPTER_NAME], adapter_weights=[float(lora_strength)]) kwargs = dict( prompt=actual_prompt, num_inference_steps=int(steps), true_cfg_scale=float(true_cfg), generator=torch.Generator(device="cpu").manual_seed(seed), ) if input_image is not None: kwargs["image"] = input_image else: kwargs.update(width=int(width), height=int(height)) if float(true_cfg) > 1.0: kwargs["negative_prompt"] = negative_prompt or "" progress(0.2, desc="Generating image") with torch.no_grad(): result = p(**kwargs) image = result.images[0] map_image = draw_map(plan, image.width, image.height) if show_map and plan else None progress(1, desc="Complete") return image, map_image, f"Complete | Seed={seed} | Steps={int(steps)} | LoRA={lora_strength}" except torch.cuda.OutOfMemoryError: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() raise gr.Error("Out of GPU memory. Try a smaller image or fewer steps.") except Exception as e: traceback.print_exc() raise gr.Error(f"Generation failed: {e}") CSS = """ .gradio-container { max-width: 1600px !important; margin: auto !important; } #prompt-box textarea { min-height: 220px !important; font-family: monospace !important; } """ with gr.Blocks(title="Pixelship 2.0", theme=gr.themes.Soft(), css=CSS) as demo: gr.Markdown("# 🎨 Pixelship 2.0\nQwen Image 2.1 + Pixelship LoRA | Manual JSON Styling | Optional spatial map") with gr.Row(): with gr.Column(scale=5): prompt = gr.Textbox(label="Prompt", lines=10, max_lines=40, elem_id="prompt-box", placeholder="Describe your image...") input_image = gr.Image(label="Image to Edit (optional)", type="pil") negative_prompt = gr.Textbox(label="Negative Prompt", lines=2) with gr.Accordion("JSON Styling / Spatial Planning", open=False): use_json = gr.Checkbox(label="Use manual JSON Styling", value=False) plan_text = gr.Code(label="JSON plan (paste your own)", language="json", lines=16) show_map = gr.Checkbox(label="Show spatial layout map", value=False) gr.Markdown("The map shows planned bounding boxes, not detected objects in the generated image.") with gr.Accordion("Generation Settings", open=True): with gr.Row(): width = gr.Slider(256, 2048, value=1024, step=64, label="Width") height = gr.Slider(256, 2048, value=1024, step=64, label="Height") with gr.Row(): steps = gr.Slider(1, 100, value=30, step=1, label="Steps") true_cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.1, label="True CFG") lora_strength = gr.Slider(0, 2, value=0.75, step=0.05, label="LoRA Strength") with gr.Row(): seed = gr.Number(label="Seed", value=42, precision=0) random_seed = gr.Checkbox(label="Random Seed", value=True) generate_button = gr.Button("Generate", variant="primary", size="lg") with gr.Column(scale=5): output_image = gr.Image(label="Generated Image", type="pil", format="png", height=600) map_output = gr.Image(label="Spatial Layout Map", type="pil", format="png") status = gr.Textbox(label="Status", interactive=False) generate_button.click( fn=generate, inputs=[prompt, input_image, negative_prompt, width, height, steps, true_cfg, lora_strength, seed, random_seed, use_json, plan_text, show_map], outputs=[output_image, map_output, status], concurrency_limit=1, ) if __name__ == "__main__": print("Starting Pixelship 2.0 Gradio server...", flush=True) demo.queue(max_size=10, default_concurrency_limit=1) demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)