Blame for app.py

"""
ArchiCAD AI Renderer — Gradio app for Hugging Face ZeroGPU Spaces.

Takes an ArchiCAD viewport export (shaded 3D view, white/clay model, elevation,
or line drawing) and turns it into a photorealistic render while preserving the
geometry. Under the hood:

  * Depth-Anything-V2 (depth map)  ─┐
  * Canny edges (linework)          ├─> multi-ControlNet SDXL (RealVisXL V5)
  * Prompt builder (scene/time/etc) ─┘        │
                                              └─> optional ControlNet img2img
                                                  "detail refine" + 1.5-2x upscale
"""

import gc
import random

import cv2
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DPMSolverMultistepScheduler,
    StableDiffusionXLControlNetImg2ImgPipeline,
    StableDiffusionXLControlNetPipeline,
)
from PIL import Image, ImageOps
from transformers import pipeline as hf_pipeline

# ----------------------------------------------------------------------------
# Model loading (done once at import; ZeroGPU moves to GPU on demand)
# ----------------------------------------------------------------------------
DTYPE = torch.float16
BASE_MODEL = "SG161222/RealVisXL_V5.0"
VAE_MODEL = "madebyollin/sdxl-vae-fp16-fix"
CN_DEPTH = "diffusers/controlnet-depth-sdxl-1.0-small"
CN_CANNY = "diffusers/controlnet-canny-sdxl-1.0-small"
DEPTH_MODEL = "depth-anything/Depth-Anything-V2-Small-hf"

print("Loading models…")
vae = AutoencoderKL.from_pretrained(VAE_MODEL, torch_dtype=DTYPE)
controlnets = [
    ControlNetModel.from_pretrained(CN_DEPTH, torch_dtype=DTYPE, use_safetensors=True),
    ControlNetModel.from_pretrained(CN_CANNY, torch_dtype=DTYPE, use_safetensors=True),
]
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    BASE_MODEL,
    controlnet=controlnets,
    vae=vae,
    torch_dtype=DTYPE,
    variant="fp16",
    use_safetensors=True,
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++"
)
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.to("cuda")

# Refine pipeline shares all weights with the main one (no extra VRAM)
refine_pipe = StableDiffusionXLControlNetImg2ImgPipeline(**pipe.components)
refine_pipe.to("cuda")

depth_estimator = hf_pipeline("depth-estimation", model=DEPTH_MODEL, device="cuda")
print("Models ready.")

# ----------------------------------------------------------------------------
# Prompt building
# ----------------------------------------------------------------------------
QUALITY = (
    "architectural photography, photorealistic, ultra detailed, sharp focus, "
    "natural lighting, realistic materials, professional DSLR, 8k"
)
NEGATIVE = (
    "cartoon, illustration, painting, sketch, lowres, blurry, deformed geometry, "
    "warped walls, distorted perspective, extra floors, extra windows, text, "
    "watermark, logo, oversaturated, cgi, plastic, noise, jpeg artifacts"
)

SCENES = {
    "Exterior – residential": "modern residential house exterior, landscaped garden, driveway",
    "Exterior – commercial / office": "contemporary commercial building exterior, plaza, street trees",
    "Exterior – urban / mixed-use": "urban mixed-use building, city street, pedestrians, cars",
    "Interior – living space": "interior design photography, living room, furniture, decor",
    "Interior – kitchen / dining": "interior design photography, kitchen and dining area, styled",
    "Interior – bedroom": "interior design photography, bedroom, bedding, soft textiles",
    "Interior – office / commercial": "modern office interior, workstations, plants, glazing",
    "Detail / facade close-up": "architectural facade detail, material junctions, shadows",
}

TIMES = {
    "Midday – clear sky": "bright midday sun, clear blue sky, crisp shadows",
    "Golden hour": "golden hour, warm low sun, long soft shadows, glowing sky",
    "Overcast / soft light": "overcast sky, soft diffuse light, muted tones",
    "Blue hour / dusk": "blue hour, dusk, glowing interior lights, deep blue sky",
    "Night": "night, illuminated interior lighting, warm lights, dark sky",
    "Misty morning": "early morning mist, soft sunrise light, atmospheric haze",
}

SEASONS = {
    "Summer": "lush green summer vegetation",
    "Spring": "spring blossoms, fresh green foliage",
    "Autumn": "autumn foliage, orange and red leaves",
    "Winter": "winter, light snow, bare trees",
    "Not specified": "",
}

STYLES = {
    "Photorealistic": "photorealistic archviz render, V-Ray look",
    "Editorial magazine": "architectural digest editorial photo, moody, high-end",
    "Minimal / Scandinavian": "minimal Scandinavian aesthetic, light wood, white walls",
    "Brutalist concrete": "brutalist architecture, raw board-formed concrete",
    "Warm / Mediterranean": "Mediterranean style, warm plaster, terracotta, olive trees",
}


def build_prompt(description, scene, time_of_day, season, style, materials):
    parts = [
        description.strip(),
        materials.strip(),
        SCENES.get(scene, ""),
        STYLES.get(style, ""),
        TIMES.get(time_of_day, ""),
        SEASONS.get(season, "") if scene.startswith("Exterior") else "",
        QUALITY,
    ]
    return ", ".join(p for p in parts if p)


# ----------------------------------------------------------------------------
# Control-image preparation
# ----------------------------------------------------------------------------
def fit_size(w, h, long_side):
    scale = long_side / max(w, h)
    nw, nh = int(round(w * scale / 8) * 8), int(round(h * scale / 8) * 8)
    return max(nw, 256), max(nh, 256)


def make_canny(img: Image.Image, low: int, high: int, thicken: bool) -> Image.Image:
    gray = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY)
    gray = cv2.GaussianBlur(gray, (3, 3), 0)
    edges = cv2.Canny(gray, low, high)
    if thicken:
        edges = cv2.dilate(edges, np.ones((2, 2), np.uint8), iterations=1)
    return Image.fromarray(np.stack([edges] * 3, axis=-1))


def make_depth(img: Image.Image) -> Image.Image:
    d = depth_estimator(img)["depth"]
    d = np.array(d).astype(np.float32)
    d = (d - d.min()) / max(d.max() - d.min(), 1e-6)
    d = (d * 255).astype(np.uint8)
    return Image.fromarray(np.stack([d] * 3, axis=-1)).resize(img.size, Image.LANCZOS)


def prepare(image, long_side, mode, low, high):
    img = ImageOps.exif_transpose(image).convert("RGB")
    w, h = fit_size(*img.size, long_side)
    img = img.resize((w, h), Image.LANCZOS)
    depth = make_depth(img)
    canny = make_canny(img, low, high, thicken=(mode == "Line drawing / elevation"))
    return img, depth, canny


MODE_DEFAULTS = {
    # (depth_scale, canny_scale)
    "3D view (shaded / clay model)": (0.65, 0.55),
    "Line drawing / elevation": (0.30, 0.90),
}


def on_mode_change(mode):
    d, c = MODE_DEFAULTS[mode]
    return d, c


# ----------------------------------------------------------------------------
# GPU functions
# ----------------------------------------------------------------------------
@spaces.GPU(duration=30)
def preview_controls(image, long_side, mode, canny_low, canny_high):
    if image is None:
        raise gr.Error("Upload an ArchiCAD export first.")
    _, depth, canny = prepare(image, long_side, mode, canny_low, canny_high)
    return [(depth, "Depth"), (canny, "Edges")]


@spaces.GPU(duration=150)
def render(
    image,
    description,
    scene,
    time_of_day,
    season,
    style,
    materials,
    negative_extra,
    mode,
    long_side,
    depth_scale,
    canny_scale,
    canny_low,
    canny_high,
    steps,
    guidance,
    num_variations,
    seed,
    do_refine,
    upscale,
    refine_strength,
    progress=gr.Progress(track_tqdm=True),
):
    if image is None:
        raise gr.Error("Upload an ArchiCAD export first.")

    seed = random.randint(0, 2**31 - 1) if seed is None or int(seed) < 0 else int(seed)
    prompt = build_prompt(description or "", scene, time_of_day, season, style, materials or "")
    negative = NEGATIVE + (", " + negative_extra.strip() if negative_extra else "")

    progress(0.05, desc="Extracting depth + edges")
    img, depth, canny = prepare(image, long_side, mode, canny_low, canny_high)
    w, h = img.size
    scales = [float(depth_scale), float(canny_scale)]

    progress(0.15, desc="Rendering")
    gen = torch.Generator("cuda").manual_seed(seed)
    out = pipe(
        prompt=prompt,
        negative_prompt=negative,
        image=[depth, canny],
        width=w,
        height=h,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        controlnet_conditioning_scale=scales,
        control_guidance_start=[0.0, 0.0],
        control_guidance_end=[0.8, 0.9],
        num_images_per_prompt=int(num_variations),
        generator=gen,
    ).images

    if do_refine:
        progress(0.7, desc="Refining details + upscaling")
        nw = int(round(w * upscale / 8) * 8)
        nh = int(round(h * upscale / 8) * 8)
        d_up = depth.resize((nw, nh), Image.LANCZOS)
        c_up = canny.resize((nw, nh), Image.LANCZOS)
        refined = []
        for i, im in enumerate(out):
            up = im.resize((nw, nh), Image.LANCZOS)
            r = refine_pipe(
                prompt=prompt,
                negative_prompt=negative,
                image=up,
                control_image=[d_up, c_up],
                strength=float(refine_strength),
                num_inference_steps=max(int(steps), 24),
                guidance_scale=float(guidance),
                controlnet_conditioning_scale=scales,
                generator=torch.Generator("cuda").manual_seed(seed + i),
            ).images[0]
            refined.append(r)
        out = refined

    gc.collect()
    torch.cuda.empty_cache()
    captions = [(im, f"seed {seed}" + (f"+{i}" if i else "")) for i, im in enumerate(out)]
    return captions, seed, prompt


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
TIPS = """
### Exporting from ArchiCAD for best results

1. **3D window → Perspective** (or Parallel). Set the *3D Window Settings → Model Display* to **Shading** or **Shading with Contours**.
2. Use **Sun/Shadow** on so volumes read clearly — depth estimation loves shadows.
3. **File → Save As → Image** (`.png`, ≥1500 px on the long side), or *Publisher → Image* for batches.
4. For **Elevations/Sections**, publish as PNG with a white background, black lines, and set *Line drawing* mode.
5. Hide furniture/labels/dimensions/zones you don't want interpreted.
6. For a "clay" look: set all surfaces to a single light grey material via *Model View Options → Override*.

### Tuning

| Goal | What to change |
|---|---|
| Geometry drifting | Raise **Depth** and **Edge** strength (0.8+) |
| Looks too CGI / stiff | Lower both strengths (0.4–0.5), raise guidance slightly |
| Too many weird details | Lower guidance (5–6), shorter prompt |
| Need a larger/sharper image | Enable **Refine + upscale** (1.5×) |
| Windows/doors missing | Lower Canny low threshold (e.g. 40/120) |

Material ideas for the prompt: `white lime render, charred timber cladding, black steel window frames, board-formed concrete, large glazing, standing seam zinc roof`.
"""

CSS = """
.gradio-container {max-width: 1280px !important}
#title h1 {font-weight: 700; letter-spacing: -0.02em}
"""

with gr.Blocks(title="ArchiCAD AI Renderer", theme=gr.themes.Soft(), css=CSS) as demo:
    gr.Markdown(
        "# 🏛️ ArchiCAD AI Renderer\n"
        "Upload a viewport export from ArchiCAD and get a photorealistic render that "
        "**keeps your geometry**. SDXL + depth/edge ControlNet on ZeroGPU.",
        elem_id="title",
    )

    with gr.Row():
        with gr.Column(scale=5):
            image = gr.Image(label="ArchiCAD export (3D view, clay model, elevation…)", type="pil", height=380)
            mode = gr.Radio(
                list(MODE_DEFAULTS.keys()), value="3D view (shaded / clay model)", label="Input type"
            )
            description = gr.Textbox(
                label="Describe the project",
                placeholder="two-storey family home with flat roof, large corner glazing, timber screen",
                lines=2,
            )
            materials = gr.Textbox(
                label="Materials",
                placeholder="white lime render, charred larch cladding, black aluminium windows, concrete paving",
            )
            with gr.Row():
                scene = gr.Dropdown(list(SCENES), value="Exterior – residential", label="Scene")
                style = gr.Dropdown(list(STYLES), value="Photorealistic", label="Style")
            with gr.Row():
                time_of_day = gr.Dropdown(list(TIMES), value="Golden hour", label="Lighting / time")
                season = gr.Dropdown(list(SEASONS), value="Summer", label="Season (exteriors)")

            with gr.Accordion("Control & generation settings", open=False):
                negative_extra = gr.Textbox(label="Extra negative prompt", placeholder="people, cars")
                long_side = gr.Slider(768, 1280, value=1024, step=64, label="Base resolution (long side)")
                with gr.Row():
                    depth_scale = gr.Slider(0, 1.2, value=0.65, step=0.05, label="Depth strength")
                    canny_scale = gr.Slider(0, 1.2, value=0.55, step=0.05, label="Edge strength")
                with gr.Row():
                    canny_low = gr.Slider(10, 200, value=70, step=5, label="Canny low")
                    canny_high = gr.Slider(50, 300, value=160, step=5, label="Canny high")
                with gr.Row():
                    steps = gr.Slider(15, 50, value=30, step=1, label="Steps")
                    guidance = gr.Slider(3, 12, value=6.5, step=0.5, label="Guidance (CFG)")
                with gr.Row():
                    num_variations = gr.Slider(1, 4, value=1, step=1, label="Variations")
                    seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")

            with gr.Accordion("Refine + upscale", open=True):
                do_refine = gr.Checkbox(value=True, label="Detail refine pass (img2img + ControlNet)")
                with gr.Row():
                    upscale = gr.Slider(1.0, 2.0, value=1.5, step=0.25, label="Upscale factor")
                    refine_strength = gr.Slider(0.15, 0.55, value=0.35, step=0.05, label="Refine strength")

            with gr.Row():
                preview_btn = gr.Button("👁️ Preview control maps", variant="secondary")
                run_btn = gr.Button("🎨 Render", variant="primary")

        with gr.Column(scale=6):
            gallery = gr.Gallery(label="Renders", columns=2, height=560, object_fit="contain", preview=True)
            with gr.Row():
                used_seed = gr.Number(label="Seed used", interactive=False)
            used_prompt = gr.Textbox(label="Final prompt", interactive=False, lines=3)
            with gr.Accordion("Control maps", open=False):
                control_gallery = gr.Gallery(columns=2, height=260, object_fit="contain", show_label=False)

    with gr.Accordion("📘 ArchiCAD export guide & tuning tips", open=False):
        gr.Markdown(TIPS)

    mode.change(on_mode_change, mode, [depth_scale, canny_scale])
    preview_btn.click(
        preview_controls,
        [image, long_side, mode, canny_low, canny_high],
        control_gallery,
    )
    run_btn.click(
        render,
        [
            image, description, scene, time_of_day, season, style, materials, negative_extra,
            mode, long_side, depth_scale, canny_scale, canny_low, canny_high,
            steps, guidance, num_variations, seed, do_refine, upscale, refine_strength,
        ],
        [gallery, used_seed, used_prompt],
    )

demo.queue(max_size=20).launch()