Update app.py

app.py

@@ -22,20 +22,26 @@ def query(payload, api_url, token):
 
 # Define the function to run when the button is clicked
 def run_lora(prompt, selected_lora_index):
-    if selected_lora_index is not None:  # Check if selected_lora_index is not None
+    if selected_lora_index is not None:
         selected_lora = loras[selected_lora_index]
         api_url = f"https://api-inference.huggingface.co/models/{selected_lora['repo']}"
         trigger_word = selected_lora["trigger_word"]
         token = os.getenv("API_TOKEN")
         payload = {"inputs": f"{prompt} {trigger_word}"}
+        
+        # API call
         image_bytes = query(payload, api_url, token)
+        
         if image_bytes:
             return Image.open(image_bytes)
     return "API Error or No LoRA selected"
 
+
 # Placeholder for gallery.select function
 def update_selection(selected):
-    return selected[0],  # Unpack the tuple and return the first element
+    if selected is None:
+        return None,
+    return selected[0],
 
 # Gradio UI
 with gr.Blocks(css="custom.css") as app: