Update app.py

app.py

@@ -6,79 +6,153 @@ 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'):
-    os.system("wget -O icon.jpg https://i.pinimg.com/564x/64/49/88/644988c59447eb00286834c2e70fdd6b.jpg")
+    print("Downloading icon...")
+    try:
+        icon_url = "https://i.pinimg.com/564x/64/49/88/644988c59447eb00286834c2e70fdd6b.jpg"
+        response = requests.get(icon_url)
+        response.raise_for_status() # Raise an exception for HTTP errors
+        with open('icon.jpg', 'wb') as f:
+            f.write(response.content)
+        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.
+
 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, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, huggingface_api_key=None, use_dev=False):
-    # Determine which API URL to use
+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):
     api_url = API_URL_DEV if use_dev else API_URL
 
-    # Check if the request is an API call by checking for the presence of the huggingface_api_key
-    is_api_call = huggingface_api_key is not None
-
-    if is_api_call:
-        # Use the environment variable for the API key in GUI mode
-        API_TOKEN = os.getenv("HF_READ_TOKEN")
-        headers = {"Authorization": f"Bearer {API_TOKEN}"}
+    # 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.")
     else:
-        # Validate the API key if it's an API call
-        if huggingface_api_key == "":
-            raise gr.Error("API key is required for API calls.")
-        headers = {"Authorization": f"Bearer {huggingface_api_key}"}
+        env_token = os.getenv("HF_READ_TOKEN")
+        if env_token and env_token.strip():
+            final_api_key = env_token.strip()
+            print("Using API key from HF_READ_TOKEN environment variable.")
+        else:
+            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.")
 
-    if prompt == "" or prompt is None:
-        return None
+    headers = {"Authorization": f"Bearer {final_api_key}"}
 
-    key = random.randint(0, 999)
+    if not prompt_text or prompt_text.strip() == "":
+        raise gr.Error("Prompt cannot be empty.")
 
-    prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
-    print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
+    key = random.randint(0, 99999) # Increased range for more uniqueness
 
-    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'\033[1mGeneration {key}:\033[0m {prompt}')
+    # 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.
+    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
+        else:
+            print(f'\033[1mGeneration {key} (no translation needed or failed, using original):\033[0m {prompt_text}')
+            prompt_to_use = prompt_text
+    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 is -1, generate a random seed and use it
     if seed == -1:
         seed = random.randint(1, 1000000000)
 
     payload = {
-        "inputs": prompt,
-        "is_negative": is_negative,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed,
-        "strength": strength
+        "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.
     }
 
-    response = requests.post(api_url, headers=headers, json=payload, timeout=timeout)
-    if response.status_code != 200:
-        print(f"Error: Failed to get image. Response status: {response.status_code}")
-        print(f"Response content: {response.text}")
-        if response.status_code == 503:
-            raise gr.Error(f"{response.status_code} : The model is being loaded")
-        raise gr.Error(f"{response.status_code}")
-    
+    # 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}")
+
+    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.")
+    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"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}")
+        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})')
+        print(f'\033[1mGeneration {key} completed!\033[0m (Prompt: {prompt_to_use})')
 
         # Save the image to a file and return the file path and seed
-        output_path = f"./output_{key}.png"
+        # Create output directory if it doesn't exist
+        os.makedirs('outputs', exist_ok=True)
+        output_path = f"./outputs/flux_output_{key}_{seed}.png"
         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.")
     except Exception as e:
-        print(f"Error when trying to open the image: {e}")
-        return None, None
+        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}")
+
 
 css = """
 #app-container {
-    max-width: 600px;
+    max-width: 700px; /* Slightly wider for better layout */
     margin-left: auto;
     margin-right: auto;
 }
@@ -86,51 +160,89 @@ css = """
     display: flex;
     align-items: center;
     justify-content: center;
+    margin-bottom: 10px; /* Add some space below title */
 }
 #title-icon {
-    width: 32px; /* Adjust the width of the icon as needed */
+    width: 40px; /* Adjusted icon size */
     height: auto;
     margin-right: 10px; /* Space between icon and title */
 }
 #title-text {
-    font-size: 24px; /* Adjust font size as needed */
+    font-size: 28px; /* Adjusted font size */
+    font-weight: bold;
+}
+.gr-input-label { /* Style labels for better visibility */
     font-weight: bold;
 }
 """
 
-with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
-    gr.HTML("""
-        <center>
-            <div id="title-container">
-                <img id="title-icon" src="icon.jpg" alt="Icon">
-                <h1 id="title-text">FLUX Capacitor</h1>
-            </div>
-        </center>
-    """)
+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.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(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")
-                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")
-                        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=1)
-                        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", value=0.7, minimum=0, maximum=1, step=0.001)
-                        seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1)
-                        huggingface_api_key = gr.Textbox(label="Hugging Face API Key (required for API calls)", placeholder="Enter your Hugging Face API Key here", type="password", elem_id="api-key")
-                        use_dev = gr.Checkbox(label="Use Dev API", value=False, elem_id="use-dev-checkbox")
+            with gr.Column(scale=2): # Prompt column takes more space
+                text_prompt = gr.Textbox(
+                    label="Prompt", 
+                    placeholder="Enter your creative vision here...", 
+                    lines=3, 
+                    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, 
+                    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.Row():
-            text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
+            text_button = gr.Button("Generate Image", variant='primary', elem_id="gen-button", scale=2)
+        
+        gr.Markdown("### Output")
         with gr.Row():
-            image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
-            seed_output = gr.Textbox(label="Seed Used", elem_id="seed-output")
+            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)
         
-        # Adjust the click function to include the API key and use_dev as inputs
-        text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, huggingface_api_key, use_dev], outputs=[image_output, seed_output])
-
-app.launch(show_api=True, share=False)
+        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
+            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=False, share=False)