Upload folder using huggingface_hub

api/index.py

@@ -33,63 +33,78 @@ def call_openai(pil_image):
     # Encode the image to base64
     image_data = base64.b64encode(buffered.getvalue()).decode('utf-8')
     
-    response = client.chat.completions.create(
-        model="gpt-4o",
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": "You are a product designer. I've attached a moodboard here. In one sentence, what do all of these elements have in common? Answer from a design language perspective, if you were telling another designer to create something similar, including any repeating colors and materials and shapes and textures"},
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": "data:image/jpeg;base64," + image_data,
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "You are a product designer. I've attached a moodboard here. In one sentence, what do all of these elements have in common? Answer from a design language perspective, if you were telling another designer to create something similar, including any repeating colors and materials and shapes and textures"},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": "data:image/jpeg;base64," + image_data,
+                            },
                         },
-                    },
-                ],
-            }
-        ],
-        max_tokens=300,
-    )
-    
-    return response.choices[0].message.content
+                    ],
+                }
+            ],
+            max_tokens=300,
+        )
+        
+        return response.choices[0].message.content
+
+    except openai.BadRequestError as e:
+        print(e)
+        print("e type")
+        print(type(e))
+        raise gr.Error(f"Please retry with a different moodboard file")
+    except Exception as e:
+        raise gr.Error("Unknown Error")
 
 def image_classifier(moodboard, starter_image, image_strength, prompt):
 
-    # Convert the numpy array to a PIL image
-    pil_image = Image.fromarray(moodboard.astype('uint8'))
-    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])
+    if moodboard is not None and starter_image is not None:
+
+        # Convert the numpy array to a PIL image
+        pil_image = Image.fromarray(moodboard.astype('uint8'))
+        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)
         
-        # Resize the image
-        starter_image_pil = starter_image_pil.resize((new_width, new_height), Image.LANCZOS)
-    
-    openai_response = call_openai(pil_image)
-    openai_response = openai_response.replace('moodboard', '')
-
-    # 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')
+        openai_response = call_openai(pil_image)
+        openai_response = openai_response.replace('moodboard', '')
+        openai_response = openai_response.replace('share', '')
+        openai_response = openai_response.replace('unified', '')
+
+        # 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')
+    else:
+        raise gr.Error(f"Please upload a moodboard to control image generation style")
         
     # Call Stable Diffusion API with the response from OpenAI
     input = {
         "width": 768,
         "height": 768,
-        "prompt": "high quality render of " + prompt + ", " + openai_response[20:],
+        "prompt": "high quality render of " + prompt + ", " + openai_response[12:],
         "negative_prompt": "worst quality, low quality, illustration, 2d, painting, cartoons, sketch",
         "refine": "expert_ensemble_refiner",
         "image": "data:image/jpeg;base64," + starter_image_base64, 
@@ -125,12 +140,5 @@ def image_classifier(moodboard, starter_image, image_strength, prompt):
     
     return [img1, img2, img3]  # Return the image object
 
-
-# app = Flask(__name__)
-# os.environ.get("REPLICATE_API_TOKEN")
-
-# @app.route("/")
-# def index():
-
 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)