Update app.py

app.py

@@ -3,47 +3,36 @@ import requests
 import io
 import random
 import os
-import time
 from PIL import Image
 import json
+from dotenv import load_dotenv
 
-# Base API URL for Hugging Face inference 
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-API_TOKEN = os.getenv("HF_READ_TOKEN")
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
-timeout = 100
-
+# Load environment variables
+load_dotenv()
 
+# Get API token from environment variable 
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable is not set")
 
 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):
-    # 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"Parameters - Steps: {steps}, CFG Scale: {cfg_scale}, Seed: {seed}, Strength: {strength}, Width: {width}, Height: {height}")
-
-    # Check if the prompt is empty or None
-    if prompt == "" or prompt is None:
-        print("Prompt is empty or None. Exiting query function.")  # Debug log
-        return None
+    
+    if not prompt:
+        raise gr.Error("Prompt cannot be empty")
 
-    # Generate a unique key for tracking the generation process
+    # Set headers with API token
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+    
+    # Generate a unique key for tracking
     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")])
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
-    print(f"Selected API token: {API_TOKEN}")  # Debug log
-
-    # Enhance the prompt with additional details for better quality
+    
+    # Enhance prompt
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'Generation {key}: {prompt}')  # Debug log
+    print(f'Generation {key}: {prompt}')
 
-    # Set the API URL based on the selected model or custom LoRA
-    if custom_lora.strip() != "":
+    # Set API URL based on model selection
+    if custom_lora.strip():
         API_URL = f"https://api-inference.huggingface.co/models/{custom_lora.strip()}"
     else:
         if model == 'Stable Diffusion XL':
@@ -245,59 +234,38 @@ def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7,
             API_URL = "https://api-inference.huggingface.co/models/artificialguybr/LogoRedmond-LogoLoraForSDXL-V2"
         if model == 'epiCPhotoGasm':
             API_URL = "https://api-inference.huggingface.co/models/Yntec/epiCPhotoGasm"
-    print(f"API URL set to: {API_URL}")  # Debug log
 
-    # Define the payload for the request
+    # Prepare payload
     payload = {
         "inputs": prompt,
-        "is_negative": is_negative,  # Whether to use a negative prompt
-        "steps": steps,  # Number of sampling steps
-        "cfg_scale": cfg_scale,  # Scale for controlling adherence to prompt
-        "seed": seed if seed != -1 else random.randint(1, 1000000000),  # Random seed for reproducibility
-        "strength": strength,  # How strongly the model should transform the image
+        "is_negative": is_negative,
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "seed": seed if seed != -1 else random.randint(1, 1000000000),
+        "strength": strength,
         "parameters": {
-            "width": width,  # Width of the generated image
-            "height": height  # Height of the generated image
+            "width": width,
+            "height": height
         }
     }
-    print(f"Payload: {json.dumps(payload, indent=2)}")  # Debug log
 
-    # Make a request to the API to generate the image
     try:
-        response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
-        print(f"Response status code: {response.status_code}")  # Debug log
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=100)
+        response.raise_for_status()
+        
+        image = Image.open(io.BytesIO(response.content))
+        print(f'Generation {key} completed successfully')
+        return image
+        
     except requests.exceptions.RequestException as e:
-        # Log any request exceptions and raise an error for the user
-        print(f"Request failed: {e}")  # Debug log
-        raise gr.Error(f"Request failed: {e}")
-
-    # Check if the response status is not successful
-    if response.status_code != 200:
-        print(f"Error: Failed to retrieve image. Response status: {response.status_code}")  # Debug log
-        print(f"Response content: {response.text}")  # Debug log
-        if response.status_code == 400:
-            raise gr.Error(f"{response.status_code}: Bad Request - There might be an issue with the input parameters.")
-        elif response.status_code == 401:
-            raise gr.Error(f"{response.status_code}: Unauthorized - Please check your API token.")
+        error_message = f"Request failed: {str(e)}"
+        if response.status_code == 401:
+            error_message = "Invalid API token. Please check your Hugging Face API token."
         elif response.status_code == 403:
-            raise gr.Error(f"{response.status_code}: Forbidden - You do not have permission to access this model.")
-        elif response.status_code == 404:
-            raise gr.Error(f"{response.status_code}: Not Found - The requested model could not be found.")
+            error_message = "Access denied. Please check your API token permissions."
         elif response.status_code == 503:
-            raise gr.Error(f"{response.status_code}: The model is being loaded. Please try again later.")
-        else:
-            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
-        image = Image.open(io.BytesIO(image_bytes))
-        print(f'Generation {key} completed! ({prompt})')  # Debug log
-        return image
-    except Exception as e:
-        # Handle any errors that occur when opening the image
-        print(f"Error while trying to open image: {e}")  # Debug log
-        return None
+            error_message = "Model is currently loading. Please try again in a few moments."
+        raise gr.Error(error_message)
 
 css = """
 footer {
@@ -307,20 +275,17 @@ footer {
 
 print("Initializing Gradio interface...")
 
-# Define the Gradio interface
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
     gr.Markdown("# AI Image Generator")
     
     with gr.Row():
         with gr.Column(scale=2):
-            # Main prompt input
             text_prompt = gr.Textbox(
                 label="Prompt",
                 placeholder="Describe what you want to create...",
                 lines=3
             )
             
-            # Negative prompt
             negative_prompt = gr.Textbox(
                 label="Negative Prompt",
                 placeholder="What should not be in the image",
@@ -328,7 +293,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                 lines=2
             )
 
-            # Custom LoRA input
             custom_lora = gr.Textbox(
                 label="Custom LoRA Path (Optional)",
                 placeholder="e.g., multimodalart/vintage-ads-flux",
@@ -336,13 +300,11 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
             )
 
         with gr.Column(scale=1):
-            # Image dimensions
             with gr.Group():
                 gr.Markdown("### Image Settings")
                 width = gr.Slider(label="Width", value=1024, minimum=512, maximum=1216, step=64)
                 height = gr.Slider(label="Height", value=1024, minimum=512, maximum=1216, step=64)
                 
-            # Generation parameters
             with gr.Group():
                 gr.Markdown("### Generation Parameters")
                 steps = gr.Slider(label="Steps", value=35, minimum=1, maximum=100, step=1)
@@ -350,7 +312,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                 strength = gr.Slider(label="Strength", value=0.7, minimum=0, maximum=1, step=0.1)
                 seed = gr.Slider(label="Seed (-1 for random)", value=-1, minimum=-1, maximum=1000000000, step=1)
 
-# Model selection
             with gr.Group():
                 gr.Markdown("### Model Selection")
                 model_search = gr.Textbox(
@@ -359,7 +320,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                     lines=1
                 )
                 
-                # Updated model list (reordered by popularity/recency)
                 models_list = [
                     "Stable Diffusion 3.5 Large",
                     "Stable Diffusion 3.5 Large Turbo",
@@ -369,7 +329,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                     "Midjourney",
                     "DreamPhotoGASM",
                     "Disney",
-                    "Leonardo AI Style Illustration",
+                    "Leonardo AI Style Illustration", 
                     "AbsoluteReality 1.8.1",
                     "Analog Redmond",
                     "Stable Diffusion 3 Medium",
@@ -379,24 +339,20 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                     "Character Design",
                     "Pixel Art XL",
                     "3D Sketchfab",
-                    "Anime Collection", # Group of anime-related models
                     "Flux Animex V2",
                     "Flux Animeo V1",
                     "Flux AestheticAnime",
                     "90s Anime Art",
                     "Softserve Anime",
-                    "Artistic Styles", # Group of artistic style models
                     "Brain Melt Acid Art",
                     "Retro Comic Flux",
                     "Purple Dreamy",
                     "SoftPasty Flux",
-                    "Specialized", # Group of specialized models
                     "Flux Logo Design",
                     "Product Design",
                     "Propaganda Poster",
                     "Movie Board",
-                    "Collage Flux",
-                    # Additional models...
+                    "Collage Flux"
                 ]
 
                 model = gr.Radio(
@@ -406,13 +362,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
                     interactive=True
                 )
 
-                def filter_models(search_term):
-                    filtered_models = [m for m in models_list if search_term.lower() in m.lower()]
-                    return gr.update(choices=filtered_models)
-
-                model_search.change(filter_models, inputs=model_search, outputs=model)
-
-    # Generate button and output
     with gr.Row():
         generate_btn = gr.Button("Generate Image", variant="primary", size="lg")
         
@@ -423,7 +372,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
             show_label=True
         )
 
-    # Set up the generation event
     generate_btn.click(
         fn=query,
         inputs=[
@@ -441,6 +389,11 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as dalle:
         outputs=image_output
     )
 
+    def filter_models(search_term):
+        filtered_models = [m for m in models_list if search_term.lower() in m.lower()]
+        return gr.update(choices=filtered_models)
+
+    model_search.change(filter_models, inputs=model_search, outputs=model)
 
-print("Launching Gradio interface...")
-dalle.launch(show_api=False, share=False)                
+if __name__ == "__main__":
+    dalle.launch(show_api=False, share=False)