Built the Sentinel Tab

dw_gateway_v2_min.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+import asyncio, json, time, uuid, io, sys, os, logging
+from pathlib import Path
+from typing import Optional, List
+import httpx
+import uvicorn
+from fastapi import FastAPI, UploadFile, File, HTTPException, Form, BackgroundTasks
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+SCRIPTS_DIR = Path(__file__).parent / "scripts"
+sys.path.insert(0, str(SCRIPTS_DIR))
+
+try:
+    from prompt_expander import expand_prompt, expand_prompt_simple, ROOM_CONTEXTS, ExpandedPrompt
+    LLM_AVAILABLE = True
+except ImportError:
+    LLM_AVAILABLE = False
+    logging.warning("prompt_expander not found — LLM expansion disabled")
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger("DreamWeaverGateway")
+COMFY = "http://127.0.0.1:8188"
+
+app = FastAPI(title="Dream Weaver API v2", version="2.0.0")
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+jobs: dict = {}
+
+async def upload_to_comfy(data: bytes, filename: str) -> str:
+    async with httpx.AsyncClient(timeout=30) as client:
+        r = await client.post(f"{COMFY}/upload/image", files={"image": (filename, data, "image/jpeg")}, data={"overwrite": "true"})
+        r.raise_for_status()
+        return r.json()["name"]
+
+def build_workflow(img_name: str, expanded: "ExpandedPrompt") -> dict:
+    return {
+        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "realvisxlV50_v50LightningBakedvae.safetensors"}},
+        "2": {"class_type": "LoadImage", "inputs": {"image": img_name, "upload": "image"}},
+        "3": {"class_type": "CLIPTextEncode", "inputs": {"text": expanded.positive_prompt, "clip": ["1", 1]}},
+        "4": {"class_type": "CLIPTextEncode", "inputs": {"text": expanded.negative_prompt, "clip": ["1", 1]}},
+        "5": {"class_type": "VAEEncode", "inputs": {"pixels": ["2", 0], "vae": ["1", 2]}},
+        "6": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["3", 0], "negative": ["4", 0], "latent_image": ["5", 0], "seed": int(time.time()) % 999983, "steps": expanded.steps, "cfg": expanded.cfg, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": expanded.denoise}},
+        "7": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["1", 2]}},
+        "8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0], "filename_prefix": f"dw_{expanded.style_name.replace(' ', '_')[:30]}"}},
+    }
+
+async def queue_prompt(workflow: dict) -> str:
+    async with httpx.AsyncClient(timeout=30) as client:
+        r = await client.post(f"{COMFY}/prompt", json={"prompt": workflow, "client_id": str(uuid.uuid4())})
+        r.raise_for_status()
+        return r.json()["prompt_id"]
+
+async def poll_result(prompt_id: str, timeout: int = 300):
+    start = time.time()
+    async with httpx.AsyncClient(timeout=10) as client:
+        while time.time() - start < timeout:
+            r = await client.get(f"{COMFY}/history/{prompt_id}")
+            if r.status_code == 200:
+                h = r.json().get(prompt_id, {})
+                imgs = [img for nd in h.get("outputs", {}).values() for img in nd.get("images", [])]
+                if imgs: return imgs[0], None
+            await asyncio.sleep(2)
+    return None, "timeout"
+
+async def background_poll(job_id: str, prompt_id: str):
+    img, err = await poll_result(prompt_id)
+    if img: jobs[job_id].update({"status": "done", "output": img, "completed": time.time()})
+    else: jobs[job_id].update({"status": "error", "error": str(err)})
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "comfyui": True, "llm_expansion": LLM_AVAILABLE, "version": "2.0.0"}
+
+@app.get("/dream-weaver/status/{job_id}")
+async def status(job_id: str):
+    job = jobs.get(job_id)
+    if not job: raise HTTPException(status_code=404, detail="Job not found")
+    res = {k: v for k, v in job.items() if k != "output"}
+    res["ready"] = job.get("status") == "done"
+    return res
+
+@app.post("/dream-weaver")
+async def dream_weaver(image: UploadFile = File(...), keywords: str = Form(default=""), room_type: str = Form(default="living_room")):
+    job_id = str(uuid.uuid4())
+    jobs[job_id] = {"status": "uploading", "created": time.time()}
+    data = await image.read()
+    comfy_name = await upload_to_comfy(data, f"dw_{job_id[:8]}.jpg")
+    kw_list = [k.strip() for k in keywords.split(",") if k.strip()]
+    expanded = await asyncio.to_thread(expand_prompt, keywords=kw_list, room_type=room_type)
+    wf = build_workflow(comfy_name, expanded)
+    prompt_id = await queue_prompt(wf)
+    jobs[job_id].update({"status": "processing", "prompt_id": prompt_id})
+    asyncio.create_task(background_poll(job_id, prompt_id))
+    return {"job_id": job_id, "status": "processing"}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8082, log_level="info")