Update app.py

app.py

@@ -4,23 +4,21 @@ import io
 import random
 import os
 import time
-from PIL import Image, UnidentifiedImageError # Added UnidentifiedImageError
+from PIL import Image, UnidentifiedImageError
 from deep_translator import GoogleTranslator
 import json
 import uuid
 from urllib.parse import quote
-import traceback # For detailed error logging
-from openai import OpenAI, RateLimitError, APIConnectionError, AuthenticationError # Added OpenAI errors
+import traceback
+from openai import OpenAI, RateLimitError, APIConnectionError, AuthenticationError
 
 # Project by Nymbo
 
 # --- Constants ---
-API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large-turbo"
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
 API_TOKEN = os.getenv("HF_READ_TOKEN")
 if not API_TOKEN:
     print("WARNING: HF_READ_TOKEN environment variable not set. API calls may fail.")
-    # Optionally, raise an error or exit if the token is essential
-    # raise ValueError("Missing required environment variable: HF_READ_TOKEN")
 headers = {"Authorization": f"Bearer {API_TOKEN}"} if API_TOKEN else {}
 timeout = 100 # seconds for API call timeout
 
@@ -42,22 +40,20 @@ absolute_image_dir = os.path.abspath(IMAGE_DIR)
 print(f"Absolute path for allowed_paths: {absolute_image_dir}")
 
 # --- OpenAI Client ---
-# Initialize OpenAI client (it will automatically look for the OPENAI_API_KEY env var)
 try:
-    # Check if the environment variable exists and is not empty
     openai_api_key = os.getenv("OPENAI_API_KEY")
     if not openai_api_key:
         print("WARNING: OPENAI_API_KEY environment variable not set or empty. Moderation will be skipped.")
         openai_client = None
     else:
-        openai_client = OpenAI() # Key is implicitly used by the library
+        openai_client = OpenAI()
         print("OpenAI client initialized successfully.")
 except Exception as e:
     print(f"ERROR: Failed to initialize OpenAI client: {e}. Moderation will be skipped.")
-    openai_client = None # Set to None so we can check later
+    openai_client = None
 
-# --- Function to query the API and return the generated image and download link ---
-def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, height=1024):
+# Function to query the API and return the generated image and download link
+def query(prompt, negative_prompt="", steps=4, cfg_scale=0, seed=-1, width=1024, height=1024):
     # Basic Input Validation
     if not prompt or not prompt.strip():
         print("WARNING: Empty prompt received.\n") # Add newline for separation
@@ -119,7 +115,7 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
             if translated_prompt != original_prompt:
                 print(f"  Translated: '{translated_prompt}'")
             print(f"  Error: {e}\n") # Add newline
-            traceback.print_exc() # Keep traceback for unexpected errors
+            traceback.print_exc()
             return None, "<p style='color: red; text-align: center;'>An unexpected error occurred during safety check. Generation blocked.</p>"
     else:
         print(f"WARNING: OpenAI client not available. Skipping moderation.")
@@ -129,14 +125,20 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
         print("") # Add newline
 
     # --- Proceed with Generation ---
+    # Add suffix (Ensure consistency with gradio1.py if desired, or keep model-specific suffixes)
     final_prompt = f"{translated_prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
 
+    # Prepare payload (Adjust parameters based on FLUX.1-schnell API requirements if different)
     payload = {
         "inputs": final_prompt,
         "parameters": {
-            "width": width, "height": height, "num_inference_steps": steps,
-            "negative_prompt": negative_prompt, "guidance_scale": cfg_scale,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": steps,
+            "guidance_scale": cfg_scale,
             "seed": seed if seed != -1 else random.randint(1, 1000000000),
+            "width": width,
+            "height": height,
+            # Add any other specific parameters for FLUX.1-schnell if needed
         }
     }
 
@@ -189,7 +191,8 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
                  print(f"  Path: '{save_path}'\n") # Add newline
                  return image, "<p style='color: red; text-align: center;'>Internal Error: Failed to confirm image file save.</p>"
 
-            space_name = "greendra-stable-diffusion-3-5-large-serverless"
+            # Use the correct space name for this Gradio app
+            space_name = "greendra-flux-1-schnell-serverless" # Updated space_name
             relative_file_url = f"gradio_api/file={save_path}"
             encoded_file_url = quote(relative_file_url)
             arinteli_url = f"{ARINTELI_REDIRECT_BASE}?download_url={encoded_file_url}&space_name={space_name}"
@@ -202,7 +205,7 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
                 f'</div>'
             )
 
-            # *** SUCCESS LOG ***
+            # *** SUCCESS LOG *** (Updated format)
             print("SUCCESS:")
             print(f"  Prompt: '{original_prompt}'")
             if translated_prompt != original_prompt:
@@ -219,7 +222,7 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
             print(f"  Path: '{save_path}'")
             print(f"  Error: {save_err}\n") # Add newline
             traceback.print_exc()
-            return image, f"<p style='color: red; text-align: center;'>Internal Error: Failed to save image file.</p>"
+            return image, f"<p style='color: red; text-align: center;'>Internal Error: Failed to save image file. Details: {save_err}</p>"
         except Exception as e:
             print("FAILED:")
             print(f"  Reason: Link Creation/Save Unexpected Error")
@@ -230,7 +233,7 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
             traceback.print_exc()
             return image, "<p style='color: red; text-align: center;'>Internal Error creating download link.</p>"
 
-    # --- Exception Handling for API Call ---
+    # --- Exception Handling for API Call --- (Updated format)
     except requests.exceptions.Timeout:
         print("FAILED:")
         print(f"  Reason: HF API Timeout")
@@ -256,15 +259,14 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
              print(f"  Translated: '{translated_prompt}'")
         print(f"  Details: '{error_text[:200]}'\n") # Add newline
 
-        # User-facing messages remain the same
+        # User-facing messages (Keep consistent with previous version of this file)
         if status_code == 503:
              estimated_time = error_data.get("estimated_time") if isinstance(error_data, dict) else None
              error_message = f"Model is loading (503), please wait." + (f" Est. time: {estimated_time:.1f}s." if estimated_time else "") + " Try again."
-        elif status_code == 400: error_message = f"Bad Request (400): Check parameters."
-        elif status_code == 422: error_message = f"Validation Error (422): Input invalid."
-        elif status_code == 401 or status_code == 403: error_message = f"Authorization Error ({status_code}): Check API Token."
-        elif status_code == 429: error_message = f"Rate Limit Error (429): Too many requests. Try again later."
-        else: error_message = f"API Error ({status_code}). Please try again."
+        elif status_code == 400: error_message = f"Bad Request (400): Check parameters. API Error: {error_text}"
+        elif status_code == 422: error_message = f"Validation Error (422): Input invalid. API Error: {error_text}"
+        elif status_code == 401 or status_code == 403: error_message = f"Authorization Error ({status_code}): Check your API Token (HF_READ_TOKEN)."
+        else: error_message = f"API Error: {status_code}. Details: {error_text}" # Adjusted generic message slightly
 
         return None, f"<p style='color: red; text-align: center;'>{error_message}</p>"
     except Exception as e:
@@ -275,10 +277,10 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
              print(f"  Translated: '{translated_prompt}'")
         print(f"  Error: {e}\n") # Add newline
         traceback.print_exc()
-        return None, f"<p style='color: red; text-align: center;'>An unexpected error occurred. Please check logs.</p>"
+        return None, f"<p style='color: red; text-align: center;'>An unexpected error occurred: {e}</p>" # Keep specific error for user
 
 
-# --- CSS Styling ---
+# CSS to style the app
 css = """
 #app-container {
     max-width: 800px;
@@ -288,34 +290,24 @@ css = """
 textarea:focus {
     background: #0d1117 !important;
 }
-#download-link-container p { /* Style the link container */
-    margin-top: 10px; /* Add some space above the link */
-    font-size: 0.9em; /* Slightly smaller text for the link message */
+#download-link-container p {
+    margin-top: 10px;
+    font-size: 0.9em;
 }
 """
 
-# --- Build the Gradio UI with Blocks ---
+# Build the Gradio UI with Blocks
 with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
+
     with gr.Column(elem_id="app-container"):
-        # --- Input Components ---
         with gr.Row():
             with gr.Column(elem_id="prompt-container"):
                 with gr.Row():
-                    text_prompt = gr.Textbox(
-                        label="Prompt",
-                        placeholder="Enter a prompt here",
-                        lines=2,
-                        elem_id="prompt-text-input"
-                    )
+                    text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=2, elem_id="prompt-text-input")
+
                 with gr.Row():
                     with gr.Accordion("Advanced Settings", open=False):
-                        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"
-                        )
+                        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():
                             width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=32)
                             height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1216, step=32)
@@ -323,28 +315,29 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
                         cfg = gr.Slider(label="CFG Scale (guidance_scale)", value=0, minimum=0, maximum=10, step=1)
                         seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1, info="Set to -1 for random seed")
 
-        # --- Action Button ---
         with gr.Row():
             text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
 
         # --- Output Components ---
         with gr.Row():
-            image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
+            image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery", show_label=False, show_download_button=False, show_share_button=False, show_fullscreen_button=False)
         with gr.Row():
+             # HTML component to display status messages or the download link
              download_link_display = gr.HTML(elem_id="download-link-container")
 
-        # --- Event Listener ---
+        # Bind the button to the query function
         text_button.click(
             query,
+            # Ensure inputs match the query function definition
             inputs=[text_prompt, negative_prompt, steps, cfg, seed, width, height],
+            # Outputs go to the image and HTML components
             outputs=[image_output, download_link_display]
         )
 
-# --- Launch the Gradio app ---
+# Launch the Gradio app with allowed_paths
 print("Starting Gradio app...")
 app.launch(
     show_api=False,
     share=False,
-    allowed_paths=[absolute_image_dir],
-    # server_name="0.0.0.0"
+    allowed_paths=[absolute_image_dir] # Added allowed_paths
 )