Update app.py

app.py

@@ -10,6 +10,7 @@ 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
 
 # Project by Nymbo
 
@@ -40,106 +41,143 @@ except OSError as e:
 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
+        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
 
 # --- 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):
-    # Renamed `strength` input as it wasn't used in the payload for txt2img
-    # Removed `sampler` input as it wasn't used in payload
-
     # Basic Input Validation
     if not prompt or not prompt.strip():
-        print("Empty prompt received.")
-        # Return None for image and an informative message for the HTML component
+        print("WARNING: Empty prompt received.\n") # Add newline for separation
         return None, "<p style='color: orange; text-align: center;'>Please enter a prompt.</p>"
 
-    key = random.randint(0, 999)
-    print(f"\n--- Generation {key} Started ---")
+    # Store original prompt for logging
+    original_prompt = prompt
 
     # Translation
     try:
-        # Using 'auto' source detection
         translated_prompt = GoogleTranslator(source='auto', target='en').translate(prompt)
     except Exception as e:
-        print(f"Translation failed: {e}. Using original prompt.")
-        translated_prompt = prompt # Fallback to original if translation fails
+        print(f"WARNING: Translation failed. Using original prompt.")
+        print(f"  Error: {e}")
+        print(f"  Prompt: '{original_prompt}'\n") # Add newline
+        translated_prompt = prompt
 
-    # Add suffix to prompt
+    # --- OpenAI Moderation Check ---
+    if openai_client:
+        try:
+            mod_response = openai_client.moderations.create(
+                model="omni-moderation-latest",
+                input=translated_prompt
+            )
+            result = mod_response.results[0]
+
+            if result.categories.sexual_minors:
+                print("BLOCKED:")
+                print(f"  Reason: sexual/minors")
+                print(f"  Prompt: '{original_prompt}'")
+                print(f"  Translated: '{translated_prompt}'\n") # Add newline
+                return None, "<p style='color: red; text-align: center;'>Prompt violates safety guidelines. Generation blocked.</p>"
+
+        except AuthenticationError:
+             print("BLOCKED:")
+             print(f"  Reason: OpenAI Auth Error")
+             print(f"  Prompt: '{original_prompt}'\n") # Add newline
+             return None, "<p style='color: red; text-align: center;'>Safety check failed. Generation blocked.</p>"
+        except RateLimitError:
+             print("BLOCKED:")
+             print(f"  Reason: OpenAI Rate Limit")
+             print(f"  Prompt: '{original_prompt}'\n") # Add newline
+             return None, "<p style='color: red; text-align: center;'>Safety check failed. Please try again later.</p>"
+        except APIConnectionError as e:
+             print("BLOCKED:")
+             print(f"  Reason: OpenAI Connection Error")
+             print(f"  Prompt: '{original_prompt}'")
+             print(f"  Error: {e}\n") # Add newline
+             return None, "<p style='color: red; text-align: center;'>Safety check failed. Please try again later.</p>"
+        except Exception as e:
+            print("BLOCKED:")
+            print(f"  Reason: OpenAI Unexpected Error")
+            print(f"  Prompt: '{original_prompt}'")
+            print(f"  Error: {e}\n") # Add newline
+            traceback.print_exc() # Keep traceback for unexpected errors
+            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.")
+        print(f"  Prompt: '{original_prompt}'\n") # Add newline
+
+    # --- Proceed with Generation ---
     final_prompt = f"{translated_prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'Generation {key} prompt: {final_prompt}')
 
-    # Prepare payload for API call
     payload = {
         "inputs": final_prompt,
-        "parameters": { # Nested parameters as per original structure
-            "width": width,
-            "height": height,
-            "num_inference_steps": steps,
-            "negative_prompt": negative_prompt,
-            "guidance_scale": cfg_scale,
+        "parameters": {
+            "width": width, "height": height, "num_inference_steps": steps,
+            "negative_prompt": negative_prompt, "guidance_scale": cfg_scale,
             "seed": seed if seed != -1 else random.randint(1, 1000000000),
         }
-        # Add other parameters here if needed (e.g., sampler if supported)
     }
 
     # API Call Section
     try:
         if not headers:
-             print("WARNING: Authorization header is missing (HF_READ_TOKEN not set?)")
-             # Handle error appropriately - maybe return an error message
+             print("FAILED:")
+             print(f"  Reason: HF Token Missing")
+             print(f"  Prompt: '{original_prompt}'\n") # Add newline
              return None, "<p style='color: red; text-align: center;'>Configuration Error: API Token missing.</p>"
 
         response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
-        response.raise_for_status() # Raises HTTPError for 4xx/5xx responses
+        response.raise_for_status()
 
         image_bytes = response.content
-        # Check for valid image data before proceeding
-        if not image_bytes or len(image_bytes) < 100: # Basic check for empty/tiny response
-             print(f"Error: Received empty or very small response content (length: {len(image_bytes)}). Potential API issue.")
+        if not image_bytes or len(image_bytes) < 100:
+             print("FAILED:")
+             print(f"  Reason: Invalid Image Data (Empty/Small)")
+             print(f"  Prompt: '{original_prompt}'")
+             print(f"  Length: {len(image_bytes)}\n") # Add newline
              return None, "<p style='color: red; text-align: center;'>API returned invalid image data.</p>"
 
         try:
             image = Image.open(io.BytesIO(image_bytes))
         except UnidentifiedImageError as img_err:
-             print(f"Error: Could not identify or open image from API response bytes: {img_err}")
-             # Optionally save the invalid bytes for debugging
-             # error_bytes_path = os.path.join(IMAGE_DIR, f"error_{key}_bytes.bin")
-             # with open(error_bytes_path, "wb") as f: f.write(image_bytes)
-             # print(f"Saved problematic bytes to {error_bytes_path}")
+             print("FAILED:")
+             print(f"  Reason: Image Processing Error")
+             print(f"  Prompt: '{original_prompt}'")
+             print(f"  Error: {img_err}\n") # Add newline
              return None, "<p style='color: red; text-align: center;'>Failed to process image data from API.</p>"
 
         # --- Save image and create download link ---
         filename = f"{int(time.time())}_{uuid.uuid4().hex[:8]}.png"
-        # save_path is relative to the script's execution directory
         save_path = os.path.join(IMAGE_DIR, filename)
-        absolute_save_path = os.path.abspath(save_path) # Get absolute path for logging
+        absolute_save_path = os.path.abspath(save_path)
 
         try:
-            # Save image explicitly as PNG
             image.save(save_path, "PNG")
 
-            # *** Verify file exists after saving ***
-            if os.path.exists(save_path):
-                file_size = os.path.getsize(save_path)
-                if file_size < 100: # Warn if the saved file is suspiciously small
-                     print(f"WARNING: Saved file {save_path} is very small ({file_size} bytes). May indicate an issue.")
-                     # Optionally return a warning message in the UI
-                     # return image, "<p style='color: orange; text-align: center;'>Warning: Saved image file is unexpectedly small.</p>"
-            else:
-                # This indicates a serious problem if save() didn't raise an error but the file isn't there
-                print(f"CRITICAL ERROR: File NOT found at {save_path} (Absolute: {absolute_save_path}) immediately after saving!")
-                return image, "<p style='color: red; text-align: center;'>Internal Error: Failed to confirm image file save.</p>"
-
-            # Get current space name from the API URL
-            space_name = "greendra-stable-diffusion-3-5-large-serverless"
-            
-            relative_file_url = f"/gradio_api/file={save_path}"
+            if not os.path.exists(save_path) or os.path.getsize(save_path) < 100:
+                 print("FAILED:")
+                 print(f"  Reason: Image Save Verification Error")
+                 print(f"  Prompt: '{original_prompt}'")
+                 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"
+            relative_file_url = f"file={save_path}"
             encoded_file_url = quote(relative_file_url)
-            # Add space_name parameter to the URL
             arinteli_url = f"{ARINTELI_REDIRECT_BASE}?download_url={encoded_file_url}&space_name={space_name}"
-            print(f"{arinteli_url}")
 
-            # Use simpler button style like the Run button
             download_html = (
                 f'<div style="text-align: center;">'
                 f'<a href="{arinteli_url}" target="_blank" class="gr-button gr-button-lg gr-button-primary">'
@@ -148,63 +186,70 @@ def query(prompt, negative_prompt, steps=4, cfg_scale=0, seed=-1, width=1024, he
                 f'</div>'
             )
 
-            print(f"--- Generation {key} Done ---")
+            # *** SUCCESS LOG ***
+            print("SUCCESS:")
+            print(f"  Prompt: '{original_prompt}'")
+            print(f"  {arinteli_url}\n") # Add newline
             return image, download_html
 
         except (OSError, IOError) as save_err:
-            # Handle errors during the file save operation
-            print(f"CRITICAL ERROR: Failed to save image to {save_path} (Absolute: {absolute_save_path}): {save_err}")
-            traceback.print_exc() # Log detailed traceback
-            return image, f"<p style='color: red; text-align: center;'>Internal Error: Failed to save image file. Details: {save_err}</p>"
+            print("FAILED:")
+            print(f"  Reason: Image Save IO Error")
+            print(f"  Prompt: '{original_prompt}'")
+            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>"
         except Exception as e:
-            # Catch any other unexpected errors during link creation/saving
-            print(f"Error during link creation or unexpected save issue: {e}")
+            print("FAILED:")
+            print(f"  Reason: Link Creation/Save Unexpected Error")
+            print(f"  Prompt: '{original_prompt}'")
+            print(f"  Error: {e}\n") # Add newline
             traceback.print_exc()
-            # Return the generated image (if available) but indicate link error
             return image, "<p style='color: red; text-align: center;'>Internal Error creating download link.</p>"
 
     # --- Exception Handling for API Call ---
     except requests.exceptions.Timeout:
-        print(f"Error: Request timed out after {timeout} seconds.")
+        print("FAILED:")
+        print(f"  Reason: HF API Timeout")
+        print(f"  Prompt: '{original_prompt}'\n") # Add newline
         return None, "<p style='color: red; text-align: center;'>Request timed out. The model is taking too long.</p>"
     except requests.exceptions.HTTPError as e:
-        # Handle HTTP errors from the API (4xx, 5xx)
         status_code = e.response.status_code
-        error_text = e.response.text # Default error text
+        error_text = e.response.text
+        error_data = {}
         try:
-            # Try to parse more specific error message from JSON response
             error_data = e.response.json()
-            error_text = error_data.get('error', error_text)
-            if isinstance(error_text, dict) and 'message' in error_text:
-                 error_text = error_text['message'] # Handle nested messages
+            parsed_error = error_data.get('error', error_text)
+            if isinstance(parsed_error, dict) and 'message' in parsed_error: error_text = parsed_error['message']
+            elif isinstance(parsed_error, list): error_text = "; ".join(map(str, parsed_error))
+            else: error_text = str(parsed_error)
         except json.JSONDecodeError:
-            pass # Keep raw text if not JSON
-
-        print(f"Error: Failed API call. Status: {status_code}, Response: {error_text}")
-
-        # Generate user-friendly messages based on status code
-        if status_code == 503: # Service Unavailable (often model loading)
-             estimated_time = error_data.get("estimated_time") if 'error_data' in locals() and isinstance(error_data, dict) else None
-             if estimated_time:
-                  error_message = f"Model is loading (503), please wait. Est. time: {estimated_time:.1f}s. Try again."
-             else:
-                  error_message = f"Service unavailable (503). Model might be loading or down. Try again later."
-        elif status_code == 400: # Bad Request (invalid parameters)
-             error_message = f"Bad Request (400): Check parameters. API Error: {error_text}"
-        elif status_code == 422: # Unprocessable Entity (validation error)
-             error_message = f"Validation Error (422): Input invalid. API Error: {error_text}"
-        elif status_code == 401 or status_code == 403: # Unauthorized / Forbidden
-             error_message = f"Authorization Error ({status_code}): Check your API Token (HF_READ_TOKEN)."
-        else: # Generic API error
-            error_message = f"API Error: {status_code}. Details: {error_text}"
-
-        # Return None for image, and the error message string for the HTML component
+            pass
+
+        print("FAILED:")
+        print(f"  Reason: HF API HTTP Error {status_code}")
+        print(f"  Prompt: '{original_prompt}'")
+        print(f"  Details: '{error_text[:200]}'\n") # Add newline
+
+        # User-facing messages remain the same
+        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."
+
         return None, f"<p style='color: red; text-align: center;'>{error_message}</p>"
     except Exception as e:
-        # Catch any other unexpected errors during the process
-        print(f"An unexpected error occurred: {e}")
-        traceback.print_exc() # Log detailed traceback
-        return None, f"<p style='color: red; text-align: center;'>An unexpected error occurred: {e}</p>"
+        print("FAILED:")
+        print(f"  Reason: Unexpected Error During Generation")
+        print(f"  Prompt: '{original_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>"
 
 
 # --- CSS Styling ---
@@ -252,11 +297,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
                             height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1216, step=32)
                         steps = gr.Slider(label="Sampling steps", value=4, minimum=1, maximum=8, step=1)
                         cfg = gr.Slider(label="CFG Scale (guidance_scale)", value=0, minimum=0, maximum=10, step=1)
-                        # Removed 'strength' slider as it wasn't used in query payload
-                        # strength = gr.Slider(label="Strength (Primarily for Img2Img)", value=0.7, minimum=0, maximum=1, step=0.001, info="Note: Strength is mainly used in Image-to-Image generation.")
                         seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1, info="Set to -1 for random seed")
-                        # Removed 'method' radio as it wasn't used in query payload
-                        # method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"], info="Note: Sampler choice might not be supported by this API.")
 
         # --- Action Button ---
         with gr.Row():
@@ -266,24 +307,20 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
         with gr.Row():
             image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
         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 click to the query function
         text_button.click(
             query,
-            # Ensure the inputs list matches the parameters of the `query` function definition
             inputs=[text_prompt, negative_prompt, steps, cfg, seed, width, height],
-            # Outputs go to the image component and the HTML component
             outputs=[image_output, download_link_display]
         )
 
 # --- Launch the Gradio app ---
 print("Starting Gradio app...")
-# Use allowed_paths with the pre-calculated absolute path to the image directory
 app.launch(
     show_api=False,
-    share=False, # Set to True only if you need a public link for direct testing
-    allowed_paths=[absolute_image_dir]
+    share=False,
+    allowed_paths=[absolute_image_dir],
+    # server_name="0.0.0.0"
 )