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("""

◈

Dream Museum

Entering the museum…

""") 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()