Update app.py

app.py

@@ -43,74 +43,80 @@ def enhance_prompt_with_qwen(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):
-    # 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
-
-    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}")
-    
     try:
-        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)
+        # 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}")
         
-        return output_path, seed
-    except Exception as e:
-        print(f"Error when trying to open the image: {e}")
-        return None, None
+        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
 
 css = """
 #app-container {