Upload folder using huggingface_hub

api/index.py

@@ -72,7 +72,7 @@ def call_openai(pil_image):
 # Could even do this 4 different times to get more diversity of renders
 # Add "simple" to prompt before word
 
-def image_classifier(moodboard, prompt):
+def image_classifier(moodboard, starter_image, image_strength, prompt):
 
     if moodboard is not None:
         pil_image = Image.fromarray(moodboard.astype('uint8'))
@@ -81,29 +81,43 @@ def image_classifier(moodboard, prompt):
 
     else:
         raise gr.Error(f"Please upload a moodboard to control image generation style")
+    
+    if starter_image is not None:
+        starter_image_pil = Image.fromarray(starter_image.astype('uint8'))
+
+        # Resize the starter image if either dimension is larger than 768 pixels
+        if starter_image_pil.size[0] > 768 or starter_image_pil.size[1] > 768:
+            # Calculate the new size while maintaining the aspect ratio
+            if starter_image_pil.size[0] > starter_image_pil.size[1]:
+                # Width is larger than height
+                new_width = 768
+                new_height = int((768 / starter_image_pil.size[0]) * starter_image_pil.size[1])
+            else:
+                # Height is larger than width
+                new_height = 768
+                new_width = int((768 / starter_image_pil.size[1]) * starter_image_pil.size[0])
+                
+                # Resize the image
+                starter_image_pil = starter_image_pil.resize((new_width, new_height), Image.LANCZOS)
+
+        # Save the starter image to a bytes buffer
+        buffered = io.BytesIO()
+        starter_image_pil.save(buffered, format="JPEG")
+        
+        # Encode the starter image to base64
+        starter_image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
 
     input = {
         "prompt": "high quality render of a " + prompt + " which " + openai_response + ", minimalist and simple mockup on a white background",
         "output_format": "jpg"
     }
     
-    try:
-        output = replicate.run(
-            "stability-ai/stable-diffusion-3",
-            input=input
-        )
-    except Exception as e:
-        raise gr.Error(f"Error: {e}")
 
-    try:
-        image_url = output[0]
-        response = requests.get(image_url)
-        img1 = Image.open(io.BytesIO(response.content))
-    except Exception as e:
-        raise gr.Error(f"Image download failed: {e}")
-
-    input["aspect_ratio"] = "3:2"
-    input["cfg"] = 6
+    if starter_image is not None:
+        input["image"] = "data:image/jpeg;base64," + starter_image_base64
+        input["prompt_strength"] = 1-image_strength
+    
     
     try:
         output = replicate.run(
@@ -128,13 +142,17 @@ def image_classifier(moodboard, prompt):
         "num_outputs": 2,
         "guidance_scale": 8.5
     }
+
+    if starter_image is not None:
+        input["image"] = "data:image/jpeg;base64," + starter_image_base64
+        input["prompt_strength"] = 1-image_strength
     
     output = replicate.run(
         "stability-ai/sdxl:7762fd07cf82c948538e41f63f77d685e02b063e37e496e96eefd46c929f9bdc",
         input=input
     )
 
-    images = [img1, img2]
+    images = [img2]
 
     for i in range(min(len(output), 2)):
         image_url = output[i]
@@ -142,13 +160,13 @@ def image_classifier(moodboard, prompt):
         images.append(Image.open(io.BytesIO(response.content)))
     
     # Add empty images if fewer than 3 were returned
-    while len(images) < 4:
+    while len(images) < 3:
         images.append(Image.new('RGB', (768, 768), 'gray'))
 
     images.reverse()
     return images
 
-demo = gr.Interface(fn=image_classifier, inputs=["image", "text"], outputs=["image", "image", "image", "image"])
+demo = gr.Interface(fn=image_classifier, inputs=["image", "image", gr.Slider(0, 1, step=0.05, value=0.2), "text"], outputs=["image", "image", "image"])
 demo.launch(share=True)