Fix some issues

app.py

@@ -5,6 +5,9 @@ from PIL import Image
 import json
 import os
 import logging
+import math 
+from tqdm import tqdm
+import time
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -31,27 +34,39 @@ def run_lora(prompt, selected_state, progress=gr.Progress(track_tqdm=True)):
     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")
+    #token = os.getenv("API_TOKEN")
     payload = {"inputs": f"{prompt} {trigger_word}"}
 
-    headers = {"Authorization": f"Bearer {token}"}
+    #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 Headers: {headers}")
     print(f"API Payload: {payload}")
-    
-    response = requests.post(api_url, headers=headers, json=payload)
-    if response.status_code == 200:
-        return Image.open(io.BytesIO(response.content))
-    else:
-        logging.error(f"API Error: {response.status_code}")
-        raise gr.Error("API Error: Unable to fetch the image.")  # Raise a Gradio error here
+
+    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
 
 
 
 with gr.Blocks(css="custom.css") as app:
-    title = gr.HTML("<h1>LoRA the Explorer</h1>")
+    title = gr.Markdown("# artificialguybr LoRA portfolio")
     selected_state = gr.State()  
     with gr.Row():
         gallery = gr.Gallery(
@@ -71,11 +86,16 @@ with gr.Blocks(css="custom.css") as app:
         update_selection,
         outputs=[prompt, selected_state]
     )
+    prompt.submit(
+        fn=run_lora,
+        inputs=[prompt, selected_state],
+        outputs=[result]
+    )
     button.click(
         fn=run_lora,
         inputs=[prompt, selected_state],
         outputs=[result]
     )
 
-app.queue(max_size=20)
-app.launch()
+app.queue(max_size=20, concurrency_count=5)
+app.launch()