ADDING CUSTOM PROVIDERS AND BYOK

app.py

@@ -17,7 +17,7 @@ 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):
+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=""):
     # Debug log to indicate function start
     print("Starting query function...")
     # Print the parameters for debugging purposes
@@ -25,6 +25,7 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
     print(f"Model: {model}")
     print(f"Custom LoRA: {custom_lora}")
     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
 
     # Check if the prompt is empty or None
     if prompt == "" or prompt is None:
@@ -35,10 +36,16 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
     key = random.randint(0, 999)
     print(f"Generated key: {key}")  # Debug log
 
-    # 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")])
+    # Determine which token to use - custom API key if provided, otherwise one of the environment variables
+    if custom_api_key.strip() != "":
+        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
+        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")
+    
     headers = {"Authorization": f"Bearer {API_TOKEN}"}
-    print(f"Selected API token: {API_TOKEN}")  # Debug log
 
     # Enhance the prompt with additional details for better quality
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
@@ -530,6 +537,16 @@ 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
+            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",
+                type="password",  # Hide the API key for security
+                elem_id="byok-input"
+            )
         with gr.Row():
             # Slider for selecting the image width
             width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=32)
@@ -615,6 +632,9 @@ 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.
 
+            ## 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.
+
             ## Width & Height
             ###### These sliders allow you to specify the resolution of your image. Default value is 1024x1024, and maximum output is 1216x1216.
 
@@ -644,8 +664,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
-    text_button.click(query, inputs=[text_prompt, model, custom_lora, negative_prompt, steps, cfg, method, seed, strength, width, height], outputs=image_output)
+    # 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)
 
 print("Launching Gradio interface...")  # Debug log
 # Launch the Gradio interface without showing the API or sharing externally