Update app.py

app.py

@@ -1,10 +1,13 @@
 import gradio as gr
 import requests
 import io
+import os
+import logging
 from PIL import Image
-import json
 from image_processing import downscale_image, limit_colors, convert_to_grayscale, convert_to_black_and_white, resize_image, DITHER_METHODS, QUANTIZATION_METHODS
-import logging
+
+# Configuração de log
+logging.basicConfig(level=logging.DEBUG)
 
 class SomeClass:
     def __init__(self):
@@ -28,44 +31,44 @@ def update_selection(selected_state: gr.SelectData):
 
 def run_lora(prompt, selected_state, progress=gr.Progress(track_tqdm=True)):
     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
+        logging.error("selected_state is None or empty. Make sure a LoRA is selected.")
+        raise gr.Error("You must select a LoRA before proceeding.")
+    
+    token = os.getenv("API_TOKEN")
+    if not token:
+        logging.error("API_TOKEN is not set.")
+        raise gr.Error("API_TOKEN is not set.")
+    
+    selected_lora_index = selected_state.index
     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}",
-        "parameters":{"negative_prompt": "bad art, ugly, watermark, deformed"},
+        "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}")
+    logging.debug(f"API Request: {api_url}")
+    logging.debug(f"API Payload: {payload}")
 
     error_count = 0
-    pbar = tqdm(total=None, desc="Loading model")
-    while(True):
-        response = requests.post(api_url, json=payload)
+    while True:
+        response = requests.post(api_url, json=payload, headers=headers)
         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)
+            logging.error(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
+            logging.error(f"Unexpected API Error: {response.status_code}")
+            raise gr.Error(f"Unexpected API Error: {response.status_code}")
+
 
 def postprocess(
     image,