add dedicated endpoint cloud mode

infer_cloud.py

@@ -59,6 +59,29 @@ def call_direct_inference_api(repo_id, token, prompt_text, generation_kwargs):
     return body
 
 
+def call_endpoint_url(endpoint_url, token, prompt_text, generation_kwargs):
+    headers = {"Content-Type": "application/json"}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+
+    payload = {
+        "inputs": prompt_text,
+        "parameters": generation_kwargs,
+        "options": {"wait_for_model": True},
+    }
+    response = requests.post(endpoint_url, headers=headers, json=payload, timeout=180)
+    try:
+        body = response.json()
+    except ValueError:
+        body = response.text
+
+    if response.status_code >= 400:
+        raise RuntimeError(f"Endpoint API error {response.status_code}: {body}")
+    if isinstance(body, dict) and body.get("error"):
+        raise RuntimeError(f"Endpoint API error: {body['error']}")
+    return body
+
+
 def run_local_fallback(args, reason):
     if not args.fallback_model_path:
         raise RuntimeError(reason)
@@ -102,7 +125,13 @@ def run_local_fallback(args, reason):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--repo-id", type=str, required=True)
+    parser.add_argument("--repo-id", type=str, default="")
+    parser.add_argument(
+        "--endpoint-url",
+        type=str,
+        default=os.getenv("HF_ENDPOINT_URL", ""),
+        help="Dedicated inference endpoint URL. Use this for true cloud inference.",
+    )
     parser.add_argument("--prompt", type=str, required=True)
     parser.add_argument("--token", type=str, default=os.getenv("HF_TOKEN"))
     parser.add_argument(
@@ -123,9 +152,10 @@ def main():
     args = parser.parse_args()
     if args.no_local_fallback:
         args.fallback_model_path = ""
+    if not args.repo_id and not args.endpoint_url:
+        raise ValueError("Pass --repo-id or --endpoint-url.")
 
     token = args.token or get_token()
-    client = InferenceClient(model=args.repo_id, token=token)
     prompt_text = build_instruction_prompt(args.prompt)
 
     generation_kwargs = {
@@ -139,26 +169,34 @@ def main():
         generation_kwargs["temperature"] = 0.01
 
     start_time = time.perf_counter()
-    try:
-        response = client.text_generation(prompt_text, **generation_kwargs)
-    except TypeError:
-        generation_kwargs.pop("return_full_text", None)
+    if args.endpoint_url:
+        try:
+            response = call_endpoint_url(args.endpoint_url, token, prompt_text, generation_kwargs)
+        except Exception as exc:
+            run_local_fallback(args, str(exc))
+            return
+    else:
+        client = InferenceClient(model=args.repo_id, token=token)
         try:
             response = client.text_generation(prompt_text, **generation_kwargs)
+        except TypeError:
+            generation_kwargs.pop("return_full_text", None)
+            try:
+                response = client.text_generation(prompt_text, **generation_kwargs)
+            except Exception as exc:
+                try:
+                    response = call_direct_inference_api(
+                        args.repo_id, token, prompt_text, generation_kwargs
+                    )
+                except Exception as direct_exc:
+                    run_local_fallback(args, f"{exc}; direct API fallback failed: {direct_exc}")
+                    return
         except Exception as exc:
             try:
-                response = call_direct_inference_api(
-                    args.repo_id, token, prompt_text, generation_kwargs
-                )
+                response = call_direct_inference_api(args.repo_id, token, prompt_text, generation_kwargs)
             except Exception as direct_exc:
                 run_local_fallback(args, f"{exc}; direct API fallback failed: {direct_exc}")
                 return
-    except Exception as exc:
-        try:
-            response = call_direct_inference_api(args.repo_id, token, prompt_text, generation_kwargs)
-        except Exception as direct_exc:
-            run_local_fallback(args, f"{exc}; direct API fallback failed: {direct_exc}")
-            return
     latency_ms = int((time.perf_counter() - start_time) * 1000)
 
     generated_text = normalize_hf_response(response).strip()