app.py

import io
import base64
from pathlib import Path

from dotenv import load_dotenv
load_dotenv()

from fastapi import Request
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image
import gradio as gr

import gallery as gallery_module
import inference as inference_module

STATIC = Path(__file__).parent / "static"


# ── Gradio generate function ──────────────────────────────────────────────────

def _gradio_generate(sketch_data, description: str, strength: float):
    if sketch_data is None or not description.strip():
        return None
    sketch_img = (
        sketch_data.get("composite")
        if isinstance(sketch_data, dict)
        else sketch_data
    )
    if sketch_img is None:
        return None
    buf = io.BytesIO()
    sketch_img.convert("RGB").save(buf, format="PNG")
    sketch_b64 = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
    image_b64 = inference_module.generate(sketch_b64, description.strip(), float(strength))
    img_bytes = base64.b64decode(image_b64)
    return Image.open(io.BytesIO(img_bytes)).convert("RGB")


# ── Gradio Blocks UI ──────────────────────────────────────────────────────────
# Gradio is required for ZeroGPU (@spaces.GPU) to work on HF Spaces.
# We redirect visitors to /museum immediately via meta-refresh + JS.

with gr.Blocks(title="Dream Museum") as demo:

    gr.HTML("""
    <meta http-equiv="refresh" content="0; url=/museum">
    <script>window.location.replace('/museum');</script>
    <style>
      footer { display: none !important; }
      @import url('https://fonts.googleapis.com/css2?family=Cormorant+Garamond:ital,wght@0,300;1,300;1,400&display=swap');
    </style>
    <div style="text-align:center;padding:40px 20px;font-family:'Cormorant Garamond',Georgia,serif">
      <p style="color:#c9a030;font-size:1.5rem;margin:0">◈</p>
      <h1 style="font-size:2.5rem;font-weight:300;font-style:italic;letter-spacing:.15em;color:#f0cc6e;margin:6px 0">Dream Museum</h1>
      <p style="color:#9880a8;font-style:italic;font-size:.95rem;margin:0">Entering the museum…</p>
    </div>
    """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            sketch_input = gr.Sketchpad(
                label="Sketch your dream",
                type="pil",
                height=440,
            )
        with gr.Column(scale=1):
            desc_input = gr.Textbox(
                label="Describe your dream",
                placeholder=(
                    "A cathedral of clouds, golden light through impossible windows, "
                    "the sensation of floating between colours…"
                ),
                lines=5,
            )
            strength_slider = gr.Slider(
                minimum=0.3, maximum=1.0, value=0.7, step=0.05,
                label="Sketch faithfulness",
                info="Low = free interpretation · High = follows your sketch",
            )
            generate_btn = gr.Button("✶  Materialize Dream", variant="primary", size="lg")
            dream_output = gr.Image(label="Your dream", type="pil", height=300)

    generate_btn.click(
        fn=_gradio_generate,
        inputs=[sketch_input, desc_input, strength_slider],
        outputs=[dream_output],
    )


# ── Custom route registration ─────────────────────────────────────────────────

def _register_custom_routes(fastapi_app):
    try:
        fastapi_app.mount(
            "/static", StaticFiles(directory=str(STATIC)), name="museum_static"
        )
    except Exception as e:
        print(f"[routes] /static mount skipped: {e}")

    @fastapi_app.get("/museum")
    async def museum():
        return FileResponse(str(STATIC / "index.html"))

    @fastapi_app.get("/galeria")
    async def galeria():
        return FileResponse(str(STATIC / "galeria.html"))

    @fastapi_app.post("/generate")
    async def generate_endpoint(request: Request):
        body        = await request.json()
        sketch_b64  = body.get("sketch_b64", "")
        description = body.get("description", "").strip()
        strength    = float(body.get("strength", 0.7))
        if not sketch_b64 or not description:
            return JSONResponse({"ok": False, "error": "Sketch and description are required"})
        try:
            image_b64 = inference_module.generate(sketch_b64, description, strength)
            return JSONResponse({"ok": True, "image": image_b64})
        except Exception as e:
            return JSONResponse({"ok": False, "error": str(e)})

    @fastapi_app.post("/gallery/public")
    async def public_gallery(request: Request):
        body   = await request.json()
        dreams = gallery_module.get_public_dreams(
            body.get("limit", 50), body.get("offset", 0)
        )
        return JSONResponse({"ok": True, "dreams": dreams})

    @fastapi_app.post("/gallery/user")
    async def user_gallery(request: Request):
        body    = await request.json()
        user_id = body.get("user_id", "")
        if not user_id:
            return JSONResponse({"ok": False, "error": "user_id required"})
        return JSONResponse({"ok": True, "dreams": gallery_module.get_user_dreams(user_id)})

    @fastapi_app.post("/gallery/save")
    async def save_dream_endpoint(request: Request):
        body    = await request.json()
        user_id = body.get("user_id", "")
        if not user_id:
            return JSONResponse({"ok": False, "error": "Login required to save dreams"})
        result = gallery_module.save_dream(
            user_id,
            body.get("sketch_b64", ""),
            body.get("image_b64", ""),
            body.get("description", ""),
            body.get("visibility", "public"),
        )
        return JSONResponse(result)

    @fastapi_app.post("/gallery/delete")
    async def delete_dream_endpoint(request: Request):
        body   = await request.json()
        result = gallery_module.delete_dream(
            body.get("dream_id", ""), body.get("user_id", "")
        )
        return JSONResponse(result)

    @fastapi_app.post("/gallery/toggle")
    async def toggle_visibility_endpoint(request: Request):
        body   = await request.json()
        result = gallery_module.toggle_visibility(
            body.get("dream_id", ""), body.get("user_id", "")
        )
        return JSONResponse(result)

    print("[routes] custom routes registered")


def _attach_routes_with_priority(fastapi_app):
    """Register custom routes, then move them to the front of the router so
    they take precedence over Gradio's own catch-all frontend routes."""
    n_before = len(fastapi_app.router.routes)
    _register_custom_routes(fastapi_app)
    new_routes = fastapi_app.router.routes[n_before:]
    del fastapi_app.router.routes[n_before:]
    fastapi_app.router.routes[0:0] = new_routes


# ── Entry point ───────────────────────────────────────────────────────────────
# demo.launch() is required for ZeroGPU (@spaces.GPU) to register correctly.
# prevent_thread_lock=True returns the real FastAPI app uvicorn is serving, so
# we can attach the museum routes to the exact object that handles requests.
# ssr_mode=False prevents Gradio 6.x from starting a Node.js SSR proxy.

if __name__ == "__main__":
    app, _local_url, _share_url = demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
        prevent_thread_lock=True,
    )

    _attach_routes_with_priority(app)

    demo.block_thread()