Update app.py

app.py

@@ -24,22 +24,18 @@ article_text = """
 """
 
 def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1, width=1024, height=1024, output_format="PNG"):
-    if prompt == "" or prompt == None:
+    if not prompt:
         return None, None, None
 
-    if lora_id.strip() == "" or lora_id == None:
+    if not lora_id.strip():
         lora_id = "black-forest-labs/FLUX.1-dev"
 
     key = random.randint(0, 999)
-
     API_URL = "https://api-inference.huggingface.co/models/" + lora_id.strip()
     
-    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
-
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
 
-    # If seed is -1, generate a random seed and use it
+    # Generate a random seed if needed
     if randomize_seed:
         seed = random.randint(1, 4294967296)
     
@@ -54,6 +50,7 @@ def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1
         }
     }
 
+    # Send the request and get the image data
     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}")
@@ -63,16 +60,17 @@ def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1
         raise gr.Error(f"{response.status_code}")
     
     try:
+        # Load and convert the image to the specified format (JPEG or PNG)
         image_bytes = response.content
         image = Image.open(io.BytesIO(image_bytes))
         print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
 
-        # Convert to specified format (JPEG or PNG)
         img_byte_arr = io.BytesIO()
         image.save(img_byte_arr, format=output_format)
         img_byte_arr.seek(0)
-        
-        return Image.open(io.BytesIO(img_byte_arr.getvalue())), seed, io.BytesIO(img_byte_arr.getvalue())
+
+        # Prepare for Gradio's download
+        return image, seed, img_byte_arr
     except Exception as e:
         print(f"Error when trying to open the image: {e}")
         return None, None, None