Update app.py

app.py

@@ -6,7 +6,7 @@ import random
 import os
 from PIL import Image
 from deep_translator import GoogleTranslator
-from gradio_client import Client  # Import the gradio client for prompt enhancement
+from gradio_client import Client  # Import the Gradio Client
 
 # os.makedirs('assets', exist_ok=True)
 if not os.path.exists('icon.jpg'):
@@ -15,108 +15,85 @@ API_URL_DEV = "https://api-inference.huggingface.co/models/black-forest-labs/FLU
 API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
 timeout = 100
 
-# Function to set the system prompt once
-def set_system_prompt():
-    client = Client("Qwen/Qwen2.5-72B-Instruct")
+def enhance_prompt(prompt):
+    """Enhance the prompt using the Mistral Nemo prompt enhancer API."""
+    client = Client("K00B404/mistral-nemo-prompt-enhancer")
     result = client.predict(
-        system="You are Qwen, an image generation prompt enhancer",
-        api_name="/modify_system_session"
+        message=prompt,
+        system_message="You are an image generation prompt enhancer and should only respond with the enhanced version of the user input image generation prompt.",
+        max_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+        api_name="/chat"
     )
-    print(f"System session modified: {result}")
     return result
 
-# Function to enhance the prompt with Qwen model
-def enhance_prompt_with_qwen(prompt):
-    client = Client("Qwen/Qwen2.5-72B-Instruct")
-    result = client.predict(
-        query=prompt,
-        history=[],
-        system="You are Qwen, an image generation prompt enhancer",
-        api_name="/model_chat"
-    )
-    
-    # Extract the relevant part of the tuple, index [0], which contains the enhanced prompt.
-    enhanced_prompt = result[0]  # This is the string we need for the image generation prompt.
-    
-    print(f"Enhanced prompt: {enhanced_prompt}")
-    return enhanced_prompt
-
-# Image generation query function
 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
+    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}"}
+    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}"}
+
+    if prompt == "" or prompt is None:
+        return None
+
+    key = random.randint(0, 999)
+
+    prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
+    print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
+
+    # Enhance the prompt using the API
+    enhanced_prompt = enhance_prompt(prompt)
+    print(f'\033[1mEnhanced Prompt:\033[0m {enhanced_prompt}')
+
+    prompt = f"{enhanced_prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+    print(f'\033[1mGeneration {key}:\033[0m {prompt}')
+
+    # 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
+    }
+
+    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}")
+    
     try:
-        # Set system prompt first
-        set_system_prompt()
-
-        # Enhance the prompt before translation
-        enhanced_prompt = enhance_prompt_with_qwen(prompt)
-
-        # Determine which API URL to use
-        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}"}
-        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}"}
-
-        if enhanced_prompt == "" or enhanced_prompt is None:
-            return None, None
-
-        key = random.randint(0, 999)
-
-        # Translate the enhanced prompt
-        enhanced_prompt = GoogleTranslator(source='ru', target='en').translate(enhanced_prompt)
-        print(f'\033[1mGeneration {key} translation:\033[0m {enhanced_prompt}')
-
-        enhanced_prompt = f"{enhanced_prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-        print(f'\033[1mGeneration {key}:\033[0m {enhanced_prompt}')
-
-        # If seed is -1, generate a random seed and use it
-        if seed == -1:
-            seed = random.randint(1, 1000000000)
-
-        payload = {
-            "inputs": enhanced_prompt,
-            "is_negative": is_negative,
-            "steps": steps,
-            "cfg_scale": cfg_scale,
-            "seed": seed,
-            "strength": strength
-        }
-
-        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}")
+        image_bytes = response.content
+        image = Image.open(io.BytesIO(image_bytes))
+        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
+
+        # Save the image to a file and return the file path and seed
+        output_path = f"./output_{key}.png"
+        image.save(output_path)
         
-        try:
-            # Attempt to open the image
-            image_bytes = response.content
-            image = Image.open(io.BytesIO(image_bytes))
-            print(f'\033[1mGeneration {key} completed!\033[0m ({enhanced_prompt})')
-
-            # Save the image to a file and return the file path and seed
-            output_path = f"./output_{key}.png"
-            image.save(output_path)
-            
-            return output_path, seed
-        except Exception as e:
-            print(f"Error when trying to open the image: {e}")
-            return None, seed  # If the image fails, return None for image, seed is still returned
-
-    except Exception as ex:
-        print(f"Error in query execution: {ex}")
-        return None, None  # If the entire process fails, return None for both
+        return output_path, seed
+    except Exception as e:
+        print(f"Error when trying to open the image: {e}")
+        return None, None
 
 css = """
 #app-container {