Update app.py

app.py

@@ -5,6 +5,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
 
 # os.makedirs('assets', exist_ok=True)
 if not os.path.exists('icon.jpg'):
@@ -13,7 +14,20 @@ 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
 
+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"
+    )
+    return result['output']  # Assuming the enhanced prompt is under 'output'
+
 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):
+    # 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
 
@@ -30,23 +44,24 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
             raise gr.Error("API key is required for API calls.")
         headers = {"Authorization": f"Bearer {huggingface_api_key}"}
 
-    if prompt == "" or prompt is None:
+    if enhanced_prompt == "" or enhanced_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}')
+    # Translate the enhanced prompt
+    enhanced_prompt = GoogleTranslator(source='ru', target='en').translate(enhanced_prompt)
+    print(f'\033[1mGeneration {key} translation:\033[0m {enhanced_prompt}')
 
-    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'\033[1mGeneration {key}:\033[0m {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": prompt,
+        "inputs": enhanced_prompt,
         "is_negative": is_negative,
         "steps": steps,
         "cfg_scale": cfg_scale,
@@ -65,7 +80,7 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     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 ({enhanced_prompt})')
 
         # Save the image to a file and return the file path and seed
         output_path = f"./output_{key}.png"