Update app.py

app.py

@@ -17,13 +17,14 @@ headers = {"Authorization": f"Bearer {API_TOKEN}"}
 # Timeout for requests
 timeout = 100
 
-def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, custom_api_key=""):
+def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, provider="hf-inference", custom_api_key=""):
     # Debug log to indicate function start
     print("Starting query function...")
     # Print the parameters for debugging purposes
     print(f"Prompt: {prompt}")
     print(f"Model: {model}")
     print(f"Custom LoRA: {custom_lora}")
+    print(f"Provider: {provider}")
     print(f"Parameters - Steps: {steps}, CFG Scale: {cfg_scale}, Seed: {seed}, Strength: {strength}, Width: {width}, Height: {height}")
     print(f"Custom API Key provided: {bool(custom_api_key.strip())}")  # Log whether a custom key was provided without printing the key
 
@@ -41,7 +42,12 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
         print("USING CUSTOM API KEY: BYOK token provided by user is being used for authentication")
         API_TOKEN = custom_api_key.strip()
     else:
-        # Randomly select an API token from available options to distribute the load
+        # If no custom key, check if provider is "hf-inference"
+        if provider != "hf-inference":
+            print(f"ERROR: Custom API key is required for {provider} provider")
+            raise gr.Error(f"A custom API key is required when using the {provider} provider. Please enter your key in the BYOK field.")
+        
+        # For hf-inference, randomly select an API token from available options to distribute the load
         API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN"), os.getenv("HF_READ_TOKEN_2"), os.getenv("HF_READ_TOKEN_3"), os.getenv("HF_READ_TOKEN_4"), os.getenv("HF_READ_TOKEN_5")])
         print("USING DEFAULT API KEY: Random environment variable token is being used for authentication")
     
@@ -275,6 +281,31 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
         if model == 'Stable Diffusion 3 Medium':
             API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3-medium-diffusers"
             prompt = f"A, {prompt}"
+            
+    # Adjust URL for different providers if needed
+    base_url = "https://api-inference.huggingface.co"
+    if provider == "fal-ai":
+        base_url = "https://api.fal.ai/hf"
+    elif provider == "nebius":
+        base_url = "https://api.nebius.ai/hf"
+    elif provider == "replicate":
+        base_url = "https://api.replicate.com/v1/models"
+        # Replicate has a different endpoint structure
+        if "https://api-inference.huggingface.co/models/" in API_URL:
+            model_id = API_URL.replace("https://api-inference.huggingface.co/models/", "")
+            API_URL = f"{base_url}/{model_id}"
+    elif provider == "together":
+        base_url = "https://api.together.xyz/inference"
+        if "https://api-inference.huggingface.co/models/" in API_URL:
+            model_id = API_URL.replace("https://api-inference.huggingface.co/models/", "")
+            API_URL = f"{base_url}/{model_id}"
+    
+    # Only update URL if we're not using a custom URL format (like for replicate)
+    if provider != "replicate" and provider != "together" and "https://api-inference.huggingface.co/models/" in API_URL:
+        model_id = API_URL.replace("https://api-inference.huggingface.co/models/", "")
+        API_URL = f"{base_url}/models/{model_id}"
+            
+    print(f"API URL set to: {API_URL}")  # Debug log
         if model == 'Duchaiten Real3D NSFW XL':
             API_URL = "https://api-inference.huggingface.co/models/stablediffusionapi/duchaiten-real3d-nsfw-xl"
         if model == 'Pixel Art XL':
@@ -351,8 +382,45 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
     }
     print(f"Payload: {json.dumps(payload, indent=2)}")  # Debug log
 
+    # Adjust payload for different providers if needed
+    if provider == "replicate" or provider == "together":
+        # These providers might have different API formats
+        if provider == "replicate":
+            payload = {
+                "version": model_id,
+                "input": {
+                    "prompt": prompt,
+                    "negative_prompt": is_negative,
+                    "num_inference_steps": steps,
+                    "guidance_scale": cfg_scale,
+                    "seed": seed if seed != -1 else random.randint(1, 1000000000),
+                    "strength": strength,
+                    "width": width,
+                    "height": height
+                }
+            }
+        elif provider == "together":
+            payload = {
+                "model": model_id,
+                "prompt": prompt,
+                "negative_prompt": is_negative,
+                "num_inference_steps": steps,
+                "guidance_scale": cfg_scale,
+                "seed": seed if seed != -1 else random.randint(1, 1000000000),
+                "strength": strength,
+                "width": width,
+                "height": height
+            }
+
     # Make a request to the API to generate the image
     try:
+        # Log which provider we're using
+        print(f"Sending request to {provider} provider")
+        
+        # Add provider to headers for HF inference if needed
+        if provider != "hf-inference" and "api-inference.huggingface.co" in API_URL:
+            headers["X-Provider"] = provider
+            
         response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
         print(f"Response status code: {response.status_code}")  # Debug log
     except requests.exceptions.RequestException as e:
@@ -378,8 +446,21 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
             raise gr.Error(f"{response.status_code}: An unexpected error occurred.")
 
     try:
-        # Attempt to read the image from the response content
-        image_bytes = response.content
+        # Handle different response formats based on provider
+        if provider == "replicate":
+            # Replicate might return a URL to the image rather than the image itself
+            result = response.json()
+            if "output" in result and isinstance(result["output"], list) and len(result["output"]) > 0:
+                image_url = result["output"][0]
+                # Fetch the image from the URL
+                image_response = requests.get(image_url, timeout=timeout)
+                image_bytes = image_response.content
+            else:
+                raise Exception(f"Unexpected Replicate response format: {result}")
+        else:
+            # Standard response with image content
+            image_bytes = response.content
+            
         image = Image.open(io.BytesIO(image_bytes))
         print(f'Generation {key} completed! ({prompt})')  # Debug log
         return image
@@ -537,16 +618,56 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme_5') as dalle:
         with gr.Row():
             # Textbox for specifying elements to exclude from the image
             negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="What should not be in the image", value="(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, misspellings, typos", lines=3, elem_id="negative-prompt-text-input")
+        
         with gr.Row():
-            # New: BYOK (Bring Your Own Key) textbox
+            # Provider selection radio buttons
+            providers_list = [
+                "hf-inference",  # Default Hugging Face Inference
+                "fal-ai",        # Fal.ai
+                "nebius",        # Nebius
+                "replicate",     # Replicate
+                "together",      # Together AI
+            ]
+            
+            provider_radio = gr.Radio(
+                choices=providers_list,
+                value="hf-inference",
+                label="Inference Provider",
+                info="Select which provider to use for image generation. Providers other than HF Inference require a custom API key.",
+                elem_id="provider-radio"
+            )
+        
+        with gr.Row():
+            # BYOK (Bring Your Own Key) textbox
             byok_textbox = gr.Textbox(
                 value="",
                 label="BYOK (Bring Your Own Key)",
-                info="Enter a custom Hugging Face API key here. When provided, this key will be used instead of the default keys.",
-                placeholder="Enter your Hugging Face API token",
+                info="Enter a custom API key here. Required for all providers except HF Inference.",
+                placeholder="Enter your API token",
                 type="password",  # Hide the API key for security
                 elem_id="byok-input"
             )
+            
+        # Function to validate provider selection based on BYOK
+        def validate_provider(api_key, provider):
+            # If no custom API key is provided, only "hf-inference" can be used
+            if not api_key.strip() and provider != "hf-inference":
+                return gr.update(value="hf-inference")
+            return gr.update(value=provider)
+            
+        # Connect the BYOK textbox to validate provider selection
+        byok_textbox.change(
+            fn=validate_provider,
+            inputs=[byok_textbox, provider_radio],
+            outputs=provider_radio
+        )
+        
+        # Also validate provider when the radio changes to ensure consistency
+        provider_radio.change(
+            fn=validate_provider,
+            inputs=[byok_textbox, provider_radio],
+            outputs=provider_radio
+        )
         with gr.Row():
             # Slider for selecting the image width
             width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=32)
@@ -632,8 +753,16 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme_5') as dalle:
             ## Negative Prompt
             ###### This box is for telling the AI what you don't want in your images. Think of it as a way to avoid certain elements. For instance, if you don't want blurry images or extra limbs showing up, this is where you'd mention it.
 
+            ## Inference Provider
+            ###### Select which AI provider to use for generating images. Different providers may have different capabilities, pricing, and performance characteristics:
+            ###### - HF Inference: Default Hugging Face inference API (uses application's default keys)
+            ###### - Fal AI: Optimized for low latency and high throughput (requires your own API key)
+            ###### - Nebius: Cloud provider with enterprise-grade infrastructure (requires your own API key)
+            ###### - Replicate: Wide variety of models with flexible deployment (requires your own API key)
+            ###### - Together: High-performance inference service (requires your own API key)
+
             ## BYOK (Bring Your Own Key)
-            ###### This allows you to use your own Hugging Face API key instead of the default keys. Enter your key here for direct access to models using your account's permissions and rate limits.
+            ###### This allows you to use your own API key instead of the default keys. Enter your key here for direct access to models using your account's permissions and rate limits. A custom key is required for all providers except HF Inference.
 
             ## Width & Height
             ###### These sliders allow you to specify the resolution of your image. Default value is 1024x1024, and maximum output is 1216x1216.
@@ -664,8 +793,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme_5') as dalle:
     with gr.Row():
         image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
 
-    # Set up button click event to call the query function with the added BYOK parameter
-    text_button.click(query, inputs=[text_prompt, model, custom_lora, negative_prompt, steps, cfg, method, seed, strength, width, height, byok_textbox], outputs=image_output)
+    # Set up button click event to call the query function with provider and BYOK parameters
+    text_button.click(query, inputs=[text_prompt, model, custom_lora, negative_prompt, steps, cfg, method, seed, strength, width, height, provider_radio, byok_textbox], outputs=image_output)
 
 print("Launching Gradio interface...")  # Debug log
 # Launch the Gradio interface without showing the API or sharing externally