Update app.py

app.py

@@ -27,17 +27,45 @@ def update_selection(selected_state: gr.SelectData):
     )
 
 def run_lora(prompt, selected_state, progress=gr.Progress(track_tqdm=True)):
-    selected_lora_index = selected_state.index
+    logging.debug(f"Inside run_lora, selected_state: {selected_state}")
+    if not selected_state:
+        logging.error("selected_state is None or empty.")
+        raise gr.Error("You must select a LoRA before proceeding.")  # Popup error when no LoRA is selected
+
+    selected_lora_index = selected_state.index  # Changed this line
     selected_lora = loras[selected_lora_index]
     api_url = f"https://api-inference.huggingface.co/models/{selected_lora['repo']}"
-    payload = {"inputs": f"{prompt} {selected_lora['trigger_word']}", "parameters": {"negative_prompt": "bad art, ugly, watermark, deformed"}}
-    response = requests.post(api_url, json=payload)
-    if response.status_code == 200:
-        original_image = Image.open(io.BytesIO(response.content))
-        processed = SomeClass()
-        processed.images = [original_image]
-        refined_image = processed.images[-1]
-        return original_image, refined_image
+    trigger_word = selected_lora["trigger_word"]
+    #token = os.getenv("API_TOKEN")
+    payload = {
+        "inputs": f"{prompt} {trigger_word}",
+        "parameters":{"negative_prompt": "bad art, ugly, watermark, deformed"},
+    }
+    #headers = {"Authorization": f"Bearer {token}"}
+    
+    # Add a print statement to display the API request
+    print(f"API Request: {api_url}")
+    #print(f"API Headers: {headers}")
+    print(f"API Payload: {payload}")
+
+    error_count = 0
+    pbar = tqdm(total=None, desc="Loading model")
+    while(True):
+        response = requests.post(api_url, json=payload)
+        if response.status_code == 200:
+            return Image.open(io.BytesIO(response.content))
+        elif response.status_code == 503:
+            #503 is triggered when the model is doing cold boot. It also gives you a time estimate from when the model is loaded but it is not super precise
+            time.sleep(1)
+            pbar.update(1)
+        elif response.status_code == 500 and error_count < 5:
+            print(response.content)
+            time.sleep(1)
+            error_count += 1
+            continue
+        else:
+            logging.error(f"API Error: {response.status_code}")
+            raise gr.Error("API Error: Unable to fetch the image.")  # Raise a Gradio error here
 
 def apply_post_processing(image, downscale, limit_colors, grayscale, black_and_white):
     processed_image = image.copy()