Update app.py

app.py

@@ -16,15 +16,12 @@ headers = {"Authorization": f"Bearer {API_TOKEN}"}
 timeout = 100
 
 # Function to query the API and return the generated image
-def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
-    if prompt == "" or prompt is None:
-        return None
+def query(prompt, negative_prompt, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
+    if not prompt:
+        raise gr.Error("Prompt cannot be empty.")
 
     key = random.randint(0, 999)
-    
-    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
-    
+
     # Translate the prompt from Russian to English if necessary
     prompt = GoogleTranslator(source='my', target='en').translate(prompt)
     print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
@@ -32,30 +29,30 @@ def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Ka
     # Add some extra flair to the prompt
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'\033[1mGeneration {key}:\033[0m {prompt}')
-    
+
     # Prepare the payload for the API call, including width and height
     payload = {
         "inputs": prompt,
-        "is_negative": is_negative,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed if seed != -1 else random.randint(1, 1000000000),
-        "strength": strength,
         "parameters": {
-            "width": width,  # Pass the width to the API
-            "height": height  # Pass the height to the API
+            "negative_prompt": negative_prompt,
+            "steps": steps,
+            "cfg_scale": cfg_scale,
+            "sampler": sampler,
+            "seed": seed if seed != -1 else random.randint(1, 1000000000),
+            "strength": strength,
+            "width": width,
+            "height": height
         }
     }
 
     # Send the request to the API and handle the response
     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}")
+        error_message = f"Error: Failed to get image. Response status: {response.status_code}"
         if response.status_code == 503:
-            raise gr.Error(f"{response.status_code} : The model is being loaded")
-        raise gr.Error(f"{response.status_code}")
-    
+            error_message += " - The model is being loaded"
+        raise gr.Error(error_message)
+
     try:
         # Convert the response content into an image
         image_bytes = response.content
@@ -63,8 +60,7 @@ def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Ka
         print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
         return image
     except Exception as e:
-        print(f"Error when trying to open the image: {e}")
-        return None
+        raise gr.Error(f"Error when trying to open the image: {e}")
 
 # CSS to style the app
 css = """
@@ -79,7 +75,7 @@ css = """
 with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
     # Add a title to the app
     gr.HTML("<center><h1>Walone AI Image Pro Large</h1></center>")
-    
+
     # Container for all the UI elements
     with gr.Column(elem_id="app-container"):
         # Add a text input for the main prompt
@@ -87,7 +83,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
             with gr.Column(elem_id="prompt-container"):
                 with gr.Row():
                     text_prompt = gr.Textbox(label="Prompt", placeholder="Prompt စာသားရေးပါ", lines=2, elem_id="prompt-text-input")
-                
+
                 # Accordion for advanced settings
                 with gr.Row():
                     with gr.Accordion("Advanced Settings", open=False):
@@ -104,13 +100,13 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
         # Add a button to trigger the image generation
         with gr.Row():
             text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
-        
+
         # Image output area to display the generated image
         with gr.Row():
             image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
-        
+
         # Bind the button to the query function with the added width and height inputs
         text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height], outputs=image_output)
 
 # Launch the Gradio app
-app.launch(show_api=False, share=False)
+app.launch(show_api=False, share=False)