Update Deployment_UI_BE.py

Deployment_UI_BE.py

@@ -789,6 +789,17 @@ async def api_middleware_infer(req: Request):
     if not internal:
         return JSONResponse({"error": "cannot resolve internal port from blob"}, status_code=400)
 
+    # ---------------- NEW: routing override for HF_TASK=text-to-image ----------------
+    env = (cspec.get("env") or [])
+    kv = {e.get("name"): e.get("value") for e in env if isinstance(e, dict) and e.get("name")}
+    hf_task = (kv.get("HF_TASK") or "").strip().lower()
+    model_id = (kv.get("MODEL_ID") or kv.get("HF_MODEL_ID") or "").strip()
+
+    if hf_task == "text-to-image" and model_id:
+        route = f"/predictions/{model_id}"
+        _INST["predictRoute"] = route
+    # -------------------------------------------------------------------------------
+
     base = f"https://{pid}-{internal}.proxy.runpod.net"
     url = f"{base}{route}"
     _log_status(f"PROMPT_ENDPOINT {url}")
@@ -799,12 +810,38 @@ async def api_middleware_infer(req: Request):
     if not isinstance(prompt, str) or not prompt.strip():
         return JSONResponse({"error": "Missing 'prompt' in request body."}, 400)
 
+    # ---------------- NEW: canonical text-to-image payload ----------------
+    if hf_task == "text-to-image":
+        body = {
+            "inputs": prompt,
+            "parameters": {
+                "num_inference_steps": 30,
+                "guidance_scale": 7.5,
+                "width": 1024,
+                "height": 1024
+            }
+        }
+        try:
+            rp = requests.post(url, json=body, timeout=120)
+            _log_status(f"PREDICT_RESP code={rp.status_code} len={len(rp.text)}")
+            if rp.ok:
+                ct = (rp.headers.get("content-type") or "").lower()
+                data = _as_json(rp) if "application/json" in ct else {"_raw": rp.text}
+                if isinstance(data, dict) and "image_b64" in data:
+                    return JSONResponse({"image_b64": data["image_b64"], "timings": data.get("timings")}, rp.status_code)
+                return JSONResponse(data, rp.status_code)
+            return JSONResponse({"error": rp.text[:400]}, status_code=rp.status_code)
+        except Exception as e:
+            _log_status(f"PREDICT_ERR {e}")
+            return JSONResponse({"error": f"inference request failed: {e}"}, status_code=502)
+    # ---------------------------------------------------------------------
+
     # HF text-classification shim
     img = (_get_container_spec().get("imageUri","")).lower()
     if "huggingface-pytorch-inference" in img and isinstance(payload.get("prompt"), str):
         payload = {"instances": [payload["prompt"]]}
 
-    # Send immediately; no readiness checks or polling gates
+    # Non-image fallback (unchanged)
     bodies = [payload, {"prompt": prompt}, {"text": prompt}, {"inputs": prompt}, {"input": prompt}]
     for body in bodies:
         try: