Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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) | |