Update app.py

app.py

@@ -10,6 +10,7 @@ import json
 import uuid
 from urllib.parse import quote
 import traceback
+from openai import OpenAI, RateLimitError, APIConnectionError, AuthenticationError
 
 # Project by Nymbo
 
@@ -38,29 +39,86 @@ except OSError as e:
 absolute_image_dir = os.path.abspath(IMAGE_DIR)
 print(f"Absolute path for allowed_paths: {absolute_image_dir}")
 
+# --- OpenAI Client ---
+try:
+    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()
+        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
+
 # 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):
-    # Removed sampler and strength as they are not explicitly used in the payload below
-    # Note: If the API endpoint *does* support sampler/strength, add them back to the 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:
         translated_prompt = GoogleTranslator(source='auto', target='en').translate(prompt)
     except Exception as e:
-        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()
+            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 ---
+    # 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."
-    print(f'Generation {key} prompt: {final_prompt}')
 
-    # Prepare the payload for the API call
+    # Prepare payload (Adjust parameters based on FLUX.1-schnell API requirements if different)
     payload = {
         "inputs": final_prompt,
         "parameters": {
@@ -70,13 +128,16 @@ def query(prompt, negative_prompt="", steps=4, cfg_scale=0, seed=-1, width=1024,
             "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
         }
     }
 
     # API Call Section
     try:
         if not headers:
-             print("WARNING: Authorization header is missing (HF_READ_TOKEN not set?)")
+             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)
@@ -84,13 +145,19 @@ def query(prompt, negative_prompt="", steps=4, cfg_scale=0, seed=-1, width=1024,
 
         image_bytes = response.content
         if not image_bytes or len(image_bytes) < 100:
-             print(f"Error: Received empty or very small response content (length: {len(image_bytes)}). Potential API issue.")
+             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}")
+             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 ---
@@ -101,86 +168,90 @@ def query(prompt, negative_prompt="", steps=4, cfg_scale=0, seed=-1, width=1024,
         try:
             image.save(save_path, "PNG")
 
-            if os.path.exists(save_path):
-                file_size = os.path.getsize(save_path)
-                if file_size < 100:
-                     print(f"WARNING: Saved file {save_path} is very small ({file_size} bytes). May indicate an issue.")
-            else:
-                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>"
-
-            # Determine space name (adjust logic if API_URL format differs)
-            try:
-                space_name = "greendra-flux-1-schnell-serverless"
-                # A more robust way might involve getting the space ID from env vars if available
-            except IndexError:
-                 print("WARNING: Could not reliably determine space name from API_URL. Using a default.")
-                 space_name = "unknown-flux-space" # Provide a fallback
-
-            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>"
 
+            # 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}"
-            print(f"{arinteli_url}")
 
-            # Use Gradio's primary button style for the link
             download_html = (
-                f'<div style="text-align: center; margin-top: 15px;">' # Added margin-top
+                f'<div style="text-align: center; margin-top: 15px;">' # Keep existing style
                 f'<a href="{arinteli_url}" target="_blank" class="gr-button gr-button-lg gr-button-primary">'
                 f'Download Image'
                 f'</a>'
                 f'</div>'
             )
 
-            print(f"--- Generation {key} Done ---")
+            # *** SUCCESS LOG *** (Updated format)
+            print("SUCCESS:")
+            print(f"  Prompt: '{original_prompt}'")
+            print(f"  URL: '{arinteli_url}'\n") # Add newline
             return image, download_html
 
         except (OSError, IOError) as save_err:
-            print(f"CRITICAL ERROR: Failed to save image to {save_path} (Absolute: {absolute_save_path}): {save_err}")
+            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. Details: {save_err}</p>"
         except Exception as e:
-            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 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(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:
         status_code = e.response.status_code
         error_text = e.response.text
+        error_data = {}
         try:
             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']
+            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
 
-        print(f"Error: Failed API call. Status: {status_code}, Response: {error_text}")
+        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 (Keep consistent with previous version of this file)
         if status_code == 503:
-             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:
-             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}"
+             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. 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:
-        print(f"An unexpected error occurred: {e}")
+        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: {e}</p>"
+        return None, f"<p style='color: red; text-align: center;'>An unexpected error occurred: {e}</p>" # Keep specific error for user
 
 
 # CSS to style the app
@@ -215,9 +286,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
                         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)
-                        steps = gr.Slider(label="Sampling steps", value=4, minimum=1, maximum=8, step=1) # Default updated based on query function default
-                        cfg = gr.Slider(label="CFG Scale (guidance_scale)", value=0, minimum=0, maximum=10, step=1) # Label updated
-                        # Removed strength and sampler sliders as they are not passed to query
+                        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)
                         seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1, info="Set to -1 for random seed")
 
         with gr.Row():