Upload folder using huggingface_hub

api/index.py

@@ -112,8 +112,7 @@ def image_classifier(moodboard, starter_image, image_strength, prompt):
         "apply_watermark": False,
         "num_inference_steps": 25,
         "prompt_strength": 1-image_strength,
-        "num_outputs": 3,
-        "disable_safety_checker": True
+        "num_outputs": 3
     }
     
     output = replicate.run(
@@ -121,26 +120,17 @@ def image_classifier(moodboard, starter_image, image_strength, prompt):
         input=input
     )
     
-    # Download the image from the URL
-    image_url = output[0]
-    print(image_url)
-    response = requests.get(image_url)
-    print(response)
-    img1 = Image.open(io.BytesIO(response.content))
-
-    image_url = output[1]
-    print(image_url)
-    response = requests.get(image_url)
-    print(response)
-    img2 = Image.open(io.BytesIO(response.content))
-
-    image_url = output[2]
-    print(image_url)
-    response = requests.get(image_url)
-    print(response)
-    img3 = Image.open(io.BytesIO(response.content))
+    images = []
+    for i in range(min(len(output), 3)):
+        image_url = output[i]
+        response = requests.get(image_url)
+        images.append(Image.open(io.BytesIO(response.content)))
     
-    return [img1, img2, img3]  # Return the image object
+    # Add empty images if fewer than 3 were returned
+    while len(images) < 3:
+        images.append(Image.new('RGB', (768, 768), 'gray'))
+
+    return images
 
 demo = gr.Interface(fn=image_classifier, inputs=["image", "image", gr.Slider(0, 1, step=0.025, value=0.2, label="Image Strength"), "text"], outputs=["image", "image", "image"])
 demo.launch(share=False)