Update xai_client.py

xai_client.py

@@ -15,7 +15,8 @@ DEFAULT_IMAGE_MODEL = "grok-imagine-image"
 DEFAULT_VIDEO_MODEL = "grok-imagine-video"
 
 IMAGE_ASPECT_RATIOS = [
-    "auto", "1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "2:1", "1:2", "19.5:9", "9:19.5", "20:9", "9:20"
+    "auto", "1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3",
+    "2:1", "1:2", "19.5:9", "9:19.5", "20:9", "9:20"
 ]
 IMAGE_RESOLUTIONS = ["1k", "2k"]
 VIDEO_ASPECT_RATIOS = ["16:9", "9:16", "1:1", "4:3", "3:4", "3:2", "2:3"]
@@ -64,12 +65,14 @@ def get_space_base_url(request: gr.Request | None) -> str | None:
     space_host = os.getenv("SPACE_HOST")
     if space_host:
         return f"https://{space_host}".rstrip("/")
+
     if request is not None:
         headers = getattr(request, "headers", None) or {}
         host = headers.get("x-forwarded-host") or headers.get("host")
         proto = headers.get("x-forwarded-proto") or "https"
         if host:
             return f"{proto}://{host}".rstrip("/")
+
     return None
 
 
@@ -84,22 +87,25 @@ def list_xai_models(api_key: str):
     headers = auth_headers(api_key)
     image_models = [DEFAULT_IMAGE_MODEL]
     video_models = [DEFAULT_VIDEO_MODEL]
+
     try:
         r = requests.get(f"{API_BASE}/image-generation-models", headers=headers, timeout=60)
         if r.ok:
             image_models = [m["id"] for m in r.json().get("models", []) if m.get("id")] or image_models
     except Exception:
         pass
+
     try:
         r = requests.get(f"{API_BASE}/video-generation-models", headers=headers, timeout=60)
         if r.ok:
             video_models = [m["id"] for m in r.json().get("models", []) if m.get("id")] or video_models
     except Exception:
         pass
+
     return (
         gr.update(choices=image_models, value=image_models[0]),
         gr.update(choices=video_models, value=video_models[0]),
-        f"Loaded {len(image_models)} image model(s) and {len(video_models)} video model(s)."
+        f"Loaded {len(image_models)} image model(s) and {len(video_models)} video model(s).",
     )
 
 
@@ -111,16 +117,19 @@ def generate_t2i(api_key, model, prompt, n, aspect_ratio, resolution, progress=g
         "n": int(n),
         "response_format": "b64_json",
     }
+
     if not payload["prompt"]:
         raise gr.Error("Please enter a prompt.")
     if aspect_ratio:
         payload["aspect_ratio"] = aspect_ratio
     if resolution:
         payload["resolution"] = resolution
+
     progress(0.2, desc="Generating images...")
     resp = requests.post(f"{API_BASE}/images/generations", headers=headers, json=payload, timeout=300)
     if not resp.ok:
         raise gr.Error(f"xAI image generation failed:\n{safe_json_error(resp)}")
+
     gallery, paths = [], []
     for i, item in enumerate(resp.json().get("data", []), start=1):
         if item.get("b64_json"):
@@ -129,35 +138,47 @@ def generate_t2i(api_key, model, prompt, n, aspect_ratio, resolution, progress=g
             out = download_url_to_temp(item["url"], ".png")
         else:
             continue
+
         paths.append(out)
         gallery.append((out, f"Image {i}"))
+
     if not paths:
         raise gr.Error("xAI returned no images.")
+
     progress(1.0, desc="Done")
     return gallery, paths[0], paths, f"Generated {len(paths)} image(s)."
 
 
 def edit_like_i2i(api_key, model, prompt, input_image_path, aspect_ratio, progress=gr.Progress(track_tqdm=False)):
     headers = auth_headers(api_key)
+
     if not input_image_path:
         raise gr.Error("Please upload an image.")
     if not (prompt or "").strip():
         raise gr.Error("Please enter a prompt.")
+
     payload = {
         "model": model or DEFAULT_IMAGE_MODEL,
         "prompt": prompt.strip(),
-        "image": {"url": file_to_data_uri(input_image_path), "type": "image_url"},
+        "image": {
+            "url": file_to_data_uri(input_image_path),
+            "type": "image_url",
+        },
         "response_format": "b64_json",
     }
+
     if aspect_ratio:
         payload["aspect_ratio"] = aspect_ratio
+
     progress(0.2, desc="Editing image...")
     resp = requests.post(f"{API_BASE}/images/edits", headers=headers, json=payload, timeout=300)
     if not resp.ok:
         raise gr.Error(f"xAI image edit failed:\n{safe_json_error(resp)}")
+
     data = resp.json().get("data", [])
     if not data:
         raise gr.Error("xAI returned no image.")
+
     item = data[0]
     if item.get("b64_json"):
         out = download_bytes_to_temp(base64.b64decode(item["b64_json"]), ".png")
@@ -165,6 +186,7 @@ def edit_like_i2i(api_key, model, prompt, input_image_path, aspect_ratio, progre
         out = download_url_to_temp(item["url"], ".png")
     else:
         raise gr.Error("xAI returned no output image payload.")
+
     progress(1.0, desc="Done")
     return out, out, f"Completed: {Path(out).name}"
 
@@ -172,29 +194,50 @@ def edit_like_i2i(api_key, model, prompt, input_image_path, aspect_ratio, progre
 def poll_video_result(api_key, request_id, timeout_seconds, poll_interval, progress):
     headers = {"Authorization": f"Bearer {api_key.strip()}"}
     started = time.time()
+
     while True:
         if time.time() - started > timeout_seconds:
             raise gr.Error("Timed out waiting for xAI video generation.")
+
         resp = requests.get(f"{API_BASE}/videos/{request_id}", headers=headers, timeout=120)
         if not resp.ok:
             raise gr.Error(f"xAI video polling failed:\n{safe_json_error(resp)}")
+
         data = resp.json()
         status = data.get("status", "unknown")
         progress(None, desc=f"Video status: {status}")
+
         if status == "done":
-            video_url = (data.get("video") or {}).get("url")
+            video = data.get("video") or {}
+            video_url = video.get("url")
             if not video_url:
                 raise gr.Error("xAI returned no video URL.")
-            return video_url, (data.get("video") or {}).get("duration"), data
-        if status == "expired":
-            raise gr.Error("xAI request expired before retrieval.")
+            return video_url, video.get("duration"), data
+
+        if status in {"failed", "error", "cancelled", "expired"}:
+            raise gr.Error(f"xAI video job ended with status: {status}\n{json.dumps(data, indent=2)}")
+
         time.sleep(int(poll_interval))
 
 
-def generate_i2v(api_key, model, prompt, uploaded_image_path, use_last_t2i_image, last_t2i_first_image, duration, aspect_ratio, resolution, timeout_seconds, poll_interval, progress=gr.Progress(track_tqdm=False)):
+def generate_i2v(
+    api_key,
+    model,
+    prompt,
+    uploaded_image_path,
+    use_last_t2i_image,
+    last_t2i_first_image,
+    duration,
+    aspect_ratio,
+    resolution,
+    timeout_seconds,
+    poll_interval,
+    progress=gr.Progress(track_tqdm=False),
+):
     source_image = last_t2i_first_image if use_last_t2i_image and last_t2i_first_image else uploaded_image_path
     if not source_image:
         raise gr.Error("Upload an image or use the first T2I result.")
+
     headers = auth_headers(api_key)
     payload = {
         "model": model or DEFAULT_VIDEO_MODEL,
@@ -203,43 +246,66 @@ def generate_i2v(api_key, model, prompt, uploaded_image_path, use_last_t2i_image
         "duration": int(duration),
         "resolution": resolution,
     }
+
     if not payload["prompt"]:
         raise gr.Error("Please enter an I2V prompt.")
     if aspect_ratio:
         payload["aspect_ratio"] = aspect_ratio
+
     progress(0.2, desc="Submitting I2V...")
     resp = requests.post(f"{API_BASE}/videos/generations", headers=headers, json=payload, timeout=300)
     if not resp.ok:
         raise gr.Error(f"xAI I2V request failed:\n{safe_json_error(resp)}")
+
     request_id = resp.json().get("request_id")
     if not request_id:
         raise gr.Error("xAI did not return request_id.")
-    video_url, actual_duration, _ = poll_video_result(api_key, request_id, int(timeout_seconds), int(poll_interval), progress)
+
+    video_url, actual_duration, _ = poll_video_result(
+        api_key, request_id, int(timeout_seconds), int(poll_interval), progress
+    )
     out = download_url_to_temp(video_url, ".mp4")
     return out, out, f"I2V complete. Request ID: {request_id}. Duration: {actual_duration}s"
 
 
-def generate_v2v(api_key, model, prompt, uploaded_video_path, timeout_seconds, poll_interval, request, progress=gr.Progress(track_tqdm=False)):
+def generate_v2v(
+    api_key,
+    model,
+    prompt,
+    uploaded_video_path,
+    timeout_seconds,
+    poll_interval,
+    request: gr.Request,
+    progress=gr.Progress(track_tqdm=False),
+):
     if not uploaded_video_path:
         raise gr.Error("Please upload an MP4 source video.")
     if Path(uploaded_video_path).suffix.lower() != ".mp4":
         raise gr.Error("xAI V2V expects .mp4 input.")
+
     headers = auth_headers(api_key)
     public_video_url = local_file_to_public_url(uploaded_video_path, request)
+
     payload = {
         "model": model or DEFAULT_VIDEO_MODEL,
         "prompt": (prompt or "").strip(),
         "video_url": public_video_url,
     }
+
     if not payload["prompt"]:
         raise gr.Error("Please enter a V2V prompt.")
+
     progress(0.2, desc="Submitting V2V...")
     resp = requests.post(f"{API_BASE}/videos/generations", headers=headers, json=payload, timeout=300)
     if not resp.ok:
         raise gr.Error(f"xAI V2V request failed:\n{safe_json_error(resp)}")
+
     request_id = resp.json().get("request_id")
     if not request_id:
         raise gr.Error("xAI did not return request_id.")
-    video_url, actual_duration, _ = poll_video_result(api_key, request_id, int(timeout_seconds), int(poll_interval), progress)
+
+    video_url, actual_duration, _ = poll_video_result(
+        api_key, request_id, int(timeout_seconds), int(poll_interval), progress
+    )
     out = download_url_to_temp(video_url, ".mp4")
-    return out, out, f"V2V complete. Request ID: {request_id}. Duration: {actual_duration}s"
+    return out, out, f"V2V complete. Request ID: {request_id}. Duration: {actual_duration}s"