Update app.py

app.py

@@ -6,148 +6,128 @@ import os
 from PIL import Image
 from deep_translator import GoogleTranslator
 
-# Create assets directory if it doesn't exist (though not strictly needed by this script anymore)
 # os.makedirs('assets', exist_ok=True)
-
-# Download icon if it doesn't exist
 if not os.path.exists('icon.jpg'):
-    print("Downloading icon...")
+    print("Downloading icon.jpg...")
     try:
-        icon_url = "https://i.pinimg.com/564x/64/49/88/644988c59447eb00286834c2e70fdd6b.jpg"
-        response = requests.get(icon_url)
+        # Use a more robust way to download, requests is already imported
+        response = requests.get("https://i.pinimg.com/564x/64/49/88/644988c59447eb00286834c2e70fdd6b.jpg", stream=True)
         response.raise_for_status() # Raise an exception for HTTP errors
         with open('icon.jpg', 'wb') as f:
-            f.write(response.content)
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
         print("Icon downloaded successfully.")
     except requests.exceptions.RequestException as e:
-        print(f"Failed to download icon.jpg: {e}. Please ensure you have internet access or place icon.jpg manually.")
-        # As a fallback, we can proceed without the icon if download fails.
-        # The gr.Image for the icon will show a broken image if 'icon.jpg' is missing.
+        print(f"Error downloading icon.jpg: {e}. Please ensure you have internet access or place icon.jpg manually.")
+        # As a fallback, you might want to skip using the icon or use a placeholder
+        # For now, the app will proceed and might show a broken image if icon.jpg is missing.
 
 API_URL_DEV = "https://lol-v2.mxflower.eu.org/api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
 API_URL = "https://lol-v2.mxflower.eu.org/api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
 timeout = 100
 
-def query(prompt_text, negative_prompt_text, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, huggingface_api_key_ui=None, use_dev=False):
+def query(prompt, negative_prompt_text, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, huggingface_api_key_ui=None, use_dev=False):
+    # Determine which API URL to use
     api_url = API_URL_DEV if use_dev else API_URL
 
-    # Determine the API token to use
-    final_api_key = None
-    if huggingface_api_key_ui and huggingface_api_key_ui.strip(): # Check if textbox has a non-empty value
-        final_api_key = huggingface_api_key_ui.strip()
-        print("Using API key from Gradio UI input.")
+    # Determine the API Token to use
+    # Priority: 1. UI input, 2. Environment variable HF_READ_TOKEN
+    auth_token = None
+    if huggingface_api_key_ui and huggingface_api_key_ui.strip(): # Check if UI key is provided and not just whitespace
+        auth_token = huggingface_api_key_ui.strip()
+        print("Using API key provided in the UI.")
     else:
-        env_token = os.getenv("HF_READ_TOKEN")
-        if env_token and env_token.strip():
-            final_api_key = env_token.strip()
+        auth_token = os.getenv("HF_READ_TOKEN")
+        if auth_token:
             print("Using API key from HF_READ_TOKEN environment variable.")
         else:
+            # If neither is available, raise an error.
             raise gr.Error("Hugging Face API Key is required. Please provide it in the 'Hugging Face API Key' field or set the HF_READ_TOKEN environment variable.")
 
-    headers = {"Authorization": f"Bearer {final_api_key}"}
+    headers = {"Authorization": f"Bearer {auth_token}"}
 
-    if not prompt_text or prompt_text.strip() == "":
-        raise gr.Error("Prompt cannot be empty.")
+    if not prompt or not prompt.strip(): # Check if prompt is None, empty, or just whitespace
+        # Optionally, return a placeholder or a message instead of None
+        # For now, returning None as per original logic for empty prompt
+        gr.Warning("Prompt cannot be empty.")
+        return None, seed # Return seed as well to match output structure
 
-    key = random.randint(0, 99999) # Increased range for more uniqueness
+    key = random.randint(0, 999)
 
-    # Translate prompt if it seems to be in Russian (basic check, can be improved)
-    # For simplicity, let's assume we always try to translate.
-    # If it's already English, GoogleTranslator often returns it as is.
+    # Translate prompt if it seems to be in Russian (simple check, can be improved)
+    # For simplicity, let's assume Russian if it contains Cyrillic characters
     try:
-        translated_prompt = GoogleTranslator(source='auto', target='en').translate(prompt_text)
-        if translated_prompt:
-            print(f'\033[1mGeneration {key} translation (auto -> en):\033[0m {translated_prompt}')
-            prompt_to_use = translated_prompt
+        # A more robust check might be needed, but this is a common heuristic
+        if any('\u0400' <= char <= '\u04FF' for char in prompt):
+            translated_prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
+            print(f'\033[1mGeneration {key} RU->EN translation:\033[0m {translated_prompt}')
+            prompt = translated_prompt
         else:
-            print(f'\033[1mGeneration {key} (no translation needed or failed, using original):\033[0m {prompt_text}')
-            prompt_to_use = prompt_text
+            print(f'\033[1mGeneration {key} using EN prompt (no translation needed).\033[0m')
     except Exception as e:
         print(f"Error during translation: {e}. Using original prompt.")
-        prompt_to_use = prompt_text
-
-
-    # Add quality enhancers
-    prompt_to_use = f"{prompt_to_use} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'\033[1mGeneration {key} (final prompt):\033[0m {prompt_to_use}')
-
-    if seed == -1:
-        seed = random.randint(1, 1000000000)
-
+        # Fallback to original prompt if translation fails
+
+    # Augment the prompt
+    augmented_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+    print(f'\033[1mGeneration {key} final prompt:\033[0m {augmented_prompt}')
+
+    # If seed is -1, generate a random seed and use it
+    current_seed = seed
+    if current_seed == -1:
+        current_seed = random.randint(1, 1000000000)
+    
+    # Note: The 'sampler' variable is passed to this function but not used in the payload.
+    # The custom API might handle sampler selection server-side or not support it.
+    # The 'is_negative' key in payload might be expecting the negative_prompt_text.
+    # Assuming the custom API expects negative prompt text under the 'is_negative' key.
+    # If it expects a boolean, this part needs adjustment.
     payload = {
-        "inputs": prompt_to_use,
-        "steps": int(steps), # Ensure steps is an int
-        "cfg_scale": float(cfg_scale), # Ensure cfg_scale is a float
-        "seed": int(seed), # Ensure seed is an int
-        "strength": float(strength) # Ensure strength is a float
-        # The 'sampler' parameter is not standard in basic HF Inference API for diffusers.
-        # It's often part of "parameters" or specific to certain model endpoints.
-        # For now, we'll omit it unless the custom proxy explicitly handles it.
-        # If "sampler" is needed, it would typically be: "parameters": {"scheduler": sampler} or similar.
+        "inputs": augmented_prompt,
+        "is_negative": negative_prompt_text, # This sends the text of negative_prompt
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "seed": current_seed,
+        "strength": strength
     }
 
-    # Add negative_prompt to payload if provided
-    if negative_prompt_text and negative_prompt_text.strip():
-        payload["negative_prompt"] = negative_prompt_text.strip()
-        print(f'\033[1mGeneration {key} (negative prompt):\033[0m {negative_prompt_text.strip()}')
-
-
-    print(f"Sending payload to {api_url}: {payload}")
+    print(f"Sending payload to {api_url}: {payload}") # For debugging
 
     try:
         response = requests.post(api_url, headers=headers, json=payload, timeout=timeout)
         response.raise_for_status() # This will raise an HTTPError for bad responses (4xx or 5xx)
     except requests.exceptions.Timeout:
-        raise gr.Error(f"Request timed out after {timeout} seconds. The model might be too busy or the request too complex.")
+        raise gr.Error(f"Request timed out after {timeout} seconds. The model might be too busy or the server is slow.")
     except requests.exceptions.HTTPError as e:
-        status_code = e.response.status_code
-        error_message = f"API Error: {status_code}."
-        try:
-            error_detail = e.response.json() # Try to get JSON error detail
-            if 'error' in error_detail:
-                error_message += f" Detail: {error_detail['error']}"
-            if 'warnings' in error_detail:
-                 error_message += f" Warnings: {error_detail['warnings']}"
-        except ValueError: # If response is not JSON
-            error_message += f" Content: {e.response.text[:200]}" # Show first 200 chars of text response
-
-        if status_code == 503: # Model loading
-            error_message = f"{status_code}: The model is currently loading. Please try again in a few moments."
-        elif status_code == 401: # Unauthorized
-             error_message = f"{status_code}: Unauthorized. Check your API Key."
-        elif status_code == 422: # Unprocessable Entity
-             error_message = f"{status_code}: Unprocessable Entity. There might be an issue with the prompt or parameters. Details: {e.response.text[:200]}"
-        
-        print(f"Error: Failed to get image. Response status: {status_code}")
+        print(f"Error: Failed to get image. Response status: {e.response.status_code}")
         print(f"Response content: {e.response.text}")
-        raise gr.Error(error_message)
-    except requests.exceptions.RequestException as e:
-        # For other network errors (DNS failure, connection refused, etc.)
-        print(f"Network error: {e}")
+        if e.response.status_code == 503:
+            raise gr.Error(f"{e.response.status_code}: Service Unavailable. The model might be loading or overloaded. Please try again later.")
+        elif e.response.status_code == 401:
+             raise gr.Error(f"{e.response.status_code}: Unauthorized. Please check your API Key.")
+        elif e.response.status_code == 400:
+             raise gr.Error(f"{e.response.status_code}: Bad Request. Please check your prompt and parameters. Details: {e.response.text[:200]}") # Show first 200 chars of error
+        else:
+            raise gr.Error(f"API Error: {e.response.status_code}. Details: {e.response.text[:200]}")
+    except requests.exceptions.RequestException as e: # Catch other network errors
         raise gr.Error(f"A network error occurred: {e}")
 
 
     try:
         image_bytes = response.content
         image = Image.open(io.BytesIO(image_bytes))
-        print(f'\033[1mGeneration {key} completed!\033[0m (Prompt: {prompt_to_use})')
+        print(f'\033[1mGeneration {key} completed!\033[0m ({augmented_prompt})')
 
         # Save the image to a file and return the file path and seed
-        # Create output directory if it doesn't exist
-        os.makedirs('outputs', exist_ok=True)
-        output_path = f"./outputs/flux_output_{key}_{seed}.png"
+        os.makedirs("outputs", exist_ok=True) # Ensure output directory exists
+        output_path = f"./outputs/output_{key}_{current_seed}.png" # Include seed in filename for uniqueness
         image.save(output_path)
         
-        return output_path, seed
-    except UnidentifiedImageError:
-        print(f"Error: The response from the API was not a valid image. Response text: {response.text[:500]}")
-        raise gr.Error("The API did not return a valid image. This might happen if the model is still loading or if there was an error with the request.")
+        return output_path, current_seed
     except Exception as e:
-        print(f"Error when trying to open or save the image: {e}")
-        # Log the raw response if it's not an image for debugging
-        if 'image_bytes' not in locals(): # if response.content was never assigned
-             print(f"Raw response was: {response.text[:500]}")
-        raise gr.Error(f"An error occurred while processing the image: {e}")
+        print(f"Error processing image response: {e}")
+        print(f"Response content that caused error: {response.content[:500]}") # Log first 500 bytes
+        raise gr.Error(f"Failed to process the image from API. The API might have returned an unexpected response. Details: {str(e)}")
 
 
 css = """
@@ -163,86 +143,107 @@ css = """
     margin-bottom: 10px; /* Add some space below title */
 }
 #title-icon {
-    width: 40px; /* Adjusted icon size */
+    width: 32px; 
     height: auto;
-    margin-right: 10px; /* Space between icon and title */
+    margin-right: 10px; 
 }
 #title-text {
-    font-size: 28px; /* Adjusted font size */
+    font-size: 24px; 
     font-weight: bold;
 }
-.gr-input-label { /* Style labels for better visibility */
-    font-weight: bold;
+.gr-box { /* Ensure accordion and other boxes have some padding */
+    padding: 10px;
 }
 """
 
-with gr.Blocks(theme='gradio/soft', css=css) as app: # Using a default theme for broader compatibility
-    with gr.Row(elem_id="title-container"):
-        if os.path.exists('icon.jpg'):
-            gr.Image(value='icon.jpg', width=40, height=40, show_label=False, interactive=False, elem_id="title-icon", container=False)
-        else:
-            gr.HTML("<span>🎨</span>", elem_id="title-icon") # Fallback if icon not found
-        gr.HTML("<h1 id='title-text'>FLUX Capacitor</h1>")
+with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
+    gr.HTML("""
+        <div style="text-align: center; margin-bottom: 20px;">
+            <div id="title-container">
+                <img id="title-icon" src="file/icon.jpg" alt="Icon"> <!-- Use file/ prefix for local files in Gradio -->
+                <h1 id="title-text">FLUX Capacitor</h1>
+            </div>
+            <p>Generate images using FLUX models. Provide your API key or ensure HF_READ_TOKEN is set.</p>
+        </div>
+    """)
 
     with gr.Column(elem_id="app-container"):
-        gr.Markdown("Generate images using FLUX.1 models via a Hugging Face Inference API endpoint.")
-        
         with gr.Row():
-            with gr.Column(scale=2): # Prompt column takes more space
+            with gr.Column(scale=3): # Give more space to prompt
                 text_prompt = gr.Textbox(
-                    label="Prompt", 
-                    placeholder="Enter your creative vision here...", 
-                    lines=3, 
+                    label="Prompt",
+                    placeholder="Enter your prompt here (Russian will be auto-translated)",
+                    lines=3, # Increased lines for prompt
                     elem_id="prompt-text-input"
                 )
                 negative_prompt = gr.Textbox(
-                    label="Negative Prompt", 
-                    placeholder="Describe what to avoid 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, 
+                    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=2, # Increased lines for negative prompt
                     elem_id="negative-prompt-text-input"
                 )
-            with gr.Column(scale=1): # Settings column
-                 with gr.Accordion("Advanced Settings & API Configuration", open=False):
-                    steps = gr.Slider(label="Sampling steps", value=30, minimum=1, maximum=100, step=1) # Adjusted default based on FLUX recommendations
-                    cfg = gr.Slider(label="CFG Scale (Guidance Scale)", value=7.0, minimum=0.0, maximum=20.0, step=0.1) # FLUX often uses lower CFG
-                    # Sampler method is often not directly controllable via basic HF Inf API unless proxy supports it
-                    # method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])
-                    strength = gr.Slider(label="Strength (for img2img/variation, less relevant for txt2img)", value=0.7, minimum=0.0, maximum=1.0, step=0.01, info="Primarily for image-to-image tasks. May have limited effect here.")
-                    seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=2147483647, step=1, info="Use -1 for a random seed.")
-                    huggingface_api_key = gr.Textbox(
-                        label="Hugging Face API Key", 
-                        placeholder="hf_xxx (Optional, uses HF_READ_TOKEN if empty)", 
-                        type="password", 
-                        elem_id="api-key"
-                    )
-                    use_dev = gr.Checkbox(label="Use FLUX.1-dev API (experimental, potentially slower)", value=False, elem_id="use-dev-checkbox")
+            with gr.Column(scale=2): # Settings column
+                 huggingface_api_key = gr.Textbox(
+                    label="Hugging Face API Key (optional)",
+                    placeholder="Uses HF_READ_TOKEN env var if empty",
+                    type="password",
+                    elem_id="api-key"
+                )
+                 use_dev = gr.Checkbox(label="Use Dev API (FLUX.1-dev)", value=False, elem_id="use-dev-checkbox")
+
+
+        with gr.Accordion("Advanced Generation Settings", open=False):
+            with gr.Row():
+                steps = gr.Slider(label="Sampling steps", value=35, minimum=1, maximum=100, step=1)
+                cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=0.5) # Allow 0.5 steps
+            with gr.Row():
+                # Sampler is not currently used in the payload. If your API supports it, add it to the payload.
+                sampler_method = gr.Radio(
+                    label="Sampling method (Note: Not sent to API)",
+                    value="DPM++ 2M Karras",
+                    choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"],
+                    # info="This setting is currently for UI only and not passed to the backend API."
+                )
+                strength = gr.Slider(label="Strength (e.g., for img2img)", value=0.7, minimum=0, maximum=1, step=0.01) # Finer steps
+            seed = gr.Slider(label="Seed (-1 for random)", value=-1, minimum=-1, maximum=2147483647, step=1) # Max 32-bit signed int
 
         with gr.Row():
-            text_button = gr.Button("Generate Image", variant='primary', elem_id="gen-button", scale=2)
+            text_button = gr.Button("Generate Image", variant='primary', elem_id="gen-button", scale=1)
         
         gr.Markdown("### Output")
         with gr.Row():
-            image_output = gr.Image(type="filepath", label="Generated Image", elem_id="gallery", height=512) # filepath is good for saved images
-            seed_output = gr.Textbox(label="Seed Used", elem_id="seed-output", interactive=False)
+            image_output = gr.Image(type="filepath", label="Generated Image", elem_id="gallery") # Use filepath for saved images
+            seed_output = gr.Textbox(label="Seed Used", interactive=False, elem_id="seed-output") # interactive=False as it's an output
         
+        # Ensure the order of inputs matches the function signature of query
         text_button.click(
-            query, 
-            inputs=[text_prompt, negative_prompt, steps, cfg, seed, strength, huggingface_api_key, use_dev], # Removed 'method' as it's not used in payload
+            query,
+            inputs=[
+                text_prompt,
+                negative_prompt, # This is passed as `negative_prompt_text`
+                steps,
+                cfg,
+                sampler_method, # Passed as `sampler`
+                seed,
+                strength,
+                huggingface_api_key, # Passed as `huggingface_api_key_ui`
+                use_dev
+            ],
             outputs=[image_output, seed_output]
         )
-        
-        gr.Markdown(
-            """
-            ---
-            *Notes:*
-            *- If the 'Hugging Face API Key' field is empty, the application will try to use the `HF_READ_TOKEN` environment variable.*
-            *- The `FLUX.1-schnell` model is used by default. Check 'Use FLUX.1-dev API' for the development version.*
-            *- Images are saved to an `outputs` subfolder in the directory where you run this script.*
-            *- Translation from any language to English is attempted for the prompt.*
-            """
-        )
 
-# Ensure the app uses show_api=False if you don't intend to expose the function as an API endpoint through Gradio
-# If you want to call this Gradio app's functions programmatically via its own API, set show_api=True
-app.launch(show_api=True, share=False)
+# To run this:
+# 1. Make sure 'gradio', 'requests', 'Pillow', 'deep_translator' are installed:
+#    pip install gradio requests Pillow deep_translator
+# 2. Optionally, set the HF_READ_TOKEN environment variable:
+#    export HF_READ_TOKEN="your_hf_api_token_here" (Linux/macOS)
+#    set HF_READ_TOKEN="your_hf_api_token_here" (Windows CMD)
+#    $env:HF_READ_TOKEN="your_hf_api_token_here" (Windows PowerShell)
+# 3. Run the script: python your_script_name.py
+
+if __name__ == "__main__":
+    # For local Gradio image serving, Gradio needs to know where the 'icon.jpg' is.
+    # If it's in the same directory, 'file/icon.jpg' should work.
+    # If you have an 'assets' folder, it would be 'file/assets/icon.jpg'.
+    app.launch(show_api=True, share=False)