Blame for app.py

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)