Upload folder using huggingface_hub

api/index.py

@@ -8,10 +8,16 @@ import requests
 import replicate
 from flask import Flask, request
 import gradio as gr
+
+import openai
 from openai import OpenAI
 
+
 from dotenv import load_dotenv, find_dotenv
 
+import json
+
+
 # Locate the .env file
 dotenv_path = find_dotenv()
 
@@ -33,79 +39,116 @@ 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 (below 20 MB in size and is of one the following formats: ['png', 'jpeg', 'gif', 'webp'])")
+    except Exception as e:
+        raise gr.Error("Unknown Error")
+
 
 def image_classifier(moodboard, prompt):
 
-    # Convert the numpy array to a PIL image
-    pil_image = Image.fromarray(moodboard.astype('uint8'))
+    if moodboard is not None:
+        pil_image = Image.fromarray(moodboard.astype('uint8'))
+
+        openai_response = call_openai(pil_image)
+        openai_response = openai_response.replace('moodboard', '')
+        openai_response = openai_response.replace('share', '')
+        openai_response = openai_response.replace('unified', '')
+    else:
+        raise gr.Error(f"Please upload a moodboard to control image generation style")
+
+    input = {
+        "prompt": "high quality render of " + prompt + ", " + openai_response[12:],
+        "negative_prompt": "worst quality, low quality, illustration, 2d, painting, cartoons, sketch",
+        "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
     
-    openai_response = call_openai(pil_image)
-    openai_response = openai_response.replace('moodboard', '')
-        
+    try:
+        output = replicate.run(
+            "stability-ai/stable-diffusion-3",
+            input=input
+        )
+        image_url = output[0]
+        response = requests.get(image_url)
+        img2 = Image.open(io.BytesIO(response.content))
+    except Exception as e:
+        raise gr.Error(f"Second image download failed: {e}")
+
     # 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",
         "apply_watermark": False,
         "num_inference_steps": 25,
-        "num_outputs": 3
+        "num_outputs": 2
     }
     
     output = replicate.run(
         "stability-ai/sdxl:7762fd07cf82c948538e41f63f77d685e02b063e37e496e96eefd46c929f9bdc",
         input=input
     )
+
+    images = [img1, img2]
+
+    for i in range(min(len(output), 2)):
+        image_url = output[i]
+        response = requests.get(image_url)
+        images.append(Image.open(io.BytesIO(response.content)))
     
-    # 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))
-    
-    return [img1, img2, img3]  # Return the image object
+    # Add empty images if fewer than 3 were returned
+    while len(images) < 4:
+        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.launch(share=True)
 
 
-# app = Flask(__name__)
-# os.environ.get("REPLICATE_API_TOKEN")
 
-# @app.route("/")
-# def index():
 
-demo = gr.Interface(fn=image_classifier, inputs=["image", "text"], outputs=["image", "image", "image"])
-demo.launch(share=True)