Update Modules/Generate_Image.py

Modules/Generate_Image.py

@@ -8,12 +8,12 @@ from typing import Annotated
 import gradio as gr
 from PIL import Image
 from huggingface_hub import InferenceClient
-from .File_System import ROOT_DIR
+from ._core import ROOT_DIR, get_hf_token, DEFAULT_PROVIDERS, handle_hf_error
 
 from app import _log_call_end, _log_call_start, _truncate_for_log
 from ._docstrings import autodoc
 
-HF_API_TOKEN = os.getenv("HF_READ_TOKEN")
+HF_API_TOKEN = get_hf_token()
 
 # Single source of truth for the LLM-facing tool description
 TOOL_SUMMARY = (
@@ -54,9 +54,8 @@ def Generate_Image(
         _log_call_end("Generate_Image", "error=empty prompt")
         raise gr.Error("Please provide a non-empty prompt.")
     enhanced_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    providers = ["auto", "replicate", "fal-ai"]
     last_error: Exception | None = None
-    for provider in providers:
+    for provider in DEFAULT_PROVIDERS:
         try:
             client = InferenceClient(api_key=HF_API_TOKEN, provider=provider)
             image = client.text_to_image(
@@ -76,21 +75,13 @@ def Generate_Image(
             
             _log_call_end("Generate_Image", f"provider={provider} size={image.size} saved_to={filename}")
             return output_path
-        except Exception as exc:  # pylint: disable=broad-except
+        except Exception as exc:
             last_error = exc
             continue
+    
     msg = str(last_error) if last_error else "Unknown error"
-    lowered = msg.lower()
-    if "404" in msg:
-        raise gr.Error(f"Model not found or unavailable: {model_id}. Check the id and your HF token access.")
-    if "503" in msg:
-        raise gr.Error("The model is warming up. Please try again shortly.")
-    if "401" in msg or "403" in msg:
-        raise gr.Error("Please duplicate the space and provide a `HF_READ_TOKEN` to enable Image and Video Generation.")
-    if ("api_key" in lowered) or ("hf auth login" in lowered) or ("unauthorized" in lowered) or ("forbidden" in lowered):
-        raise gr.Error("Please duplicate the space and provide a `HF_READ_TOKEN` to enable Image and Video Generation.")
     _log_call_end("Generate_Image", f"error={_truncate_for_log(msg, 200)}")
-    raise gr.Error(f"Image generation failed: {msg}")
+    handle_hf_error(msg, model_id, context="Image generation")
 
 
 def build_interface() -> gr.Interface: