added documentation and trying redis

app.py

@@ -7,23 +7,72 @@ import os
 import base_generator
 from openai    import OpenAI
 import base64
+import redis
+import uuid
 
 app = FastAPI()
-client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-
-
+openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+#This is a temp redis instance, replace with one on aws when going into production!
+r = redis.Redis(
+  host='cute-wildcat-41430.upstash.io',
+  port=6379,
+  password='AaHWAAIjcDFhZDVlOGUyMDQ0ZGQ0MTZmODA4ZjdkNzc4ZDVhZjUzZHAxMA',
+  ssl=True
+)
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
 
+@app.post("/virtualTryOn", summary="Virtual try on single call", 
+    description="Virtual try on single call for complete outfit try on")
+async def virtual_try_on(top_garment_image: UploadFile = File(...), 
+                        bottom_garment_image: UploadFile = File(...),
+                        top_back_garment_image: UploadFile = File(...), 
+                        bottom_back_garment_image: UploadFile = File(...),
+                        talent_trigger_word: str,
+                        talent_lora_url: str,
+                        num_images: int = 4):
+    """
+    Pass in the top garment image, bottom garment image, top back garment image, bottom back garment image, talent trigger word, talent lora url and number of images to generate.
+    Only one front garment image is required, the rest are optional.
+    Default number of images is 4.
+    """
+
+    request_id = str(uuid.uuid4())
+    r.set(request_id, "pending")
+    # step 1, make the base image
+    r.set(request_id, "Creating Base Images")
+    # step 2, combine the garment images
+    r.set(request_id, "Combining Garment Images")
+    # step 3, style the talent
+    r.set(request_id, "Styling Talent")
+    # step 4, create variations
+    r.set(request_id, "Creating Variations")
+    #step 5. Upscale the images
+    r.set(request_id, "Upscaling Images")
+    # step 6, return the result
+    r.set(request_id, "Completed")
+    return(request_id)
 
+@app.get("/status/{request_id}")
+async def check_status(request_id: str):
+    status = r.get(request_id)
+    if not status:
+        return {"error": "Invalid request_id"}
+    return {"request_id": request_id, "status": status}
 
 #endpoints related to base image generation
 
-@app.get("/makeBaseImage")
-async def make_base_image(character_lora: str, character_keyword: str, outfit_desc: str, pose_id: int   ):
+@app.get("/makeBaseImage", summary="Make a base image of the character,before adding the outfit",
+    description="Make one or more base image(s) of the character,before adding the outfit. supports multiple poses by passing in a list of pose ids")
+async def make_base_image(character_lora: str, character_keyword: str, outfit_desc: str, pose_id: int , num_outputs: int = 1):
+    """
+    Pass in the character lora, character keyword, outfit description and pose id to make a base image of the character,before adding the outfit. 
+    supports multiple poses by passing a number of outputs.
+    """
     print("make base image")
-    result = base_generator.create_image(character_lora, character_keyword, outfit_desc, pose_id)
+    result = base_generator.create_image(character_lora, character_keyword, outfit_desc, pose_id, num_outputs   )
     return(result)
 
 # Function related to virtual outfit try on
@@ -40,7 +89,7 @@ def combine_garment_images(Upper_garment: UploadFile = File(...), Lower_garment:
 # Endpoints related to virtual outfit try on
 # HOW DO WE DICESERN FRONT VS BACK IMAGES OF GARMENTS?
 
-@app.post("/styleTalent")
+@app.post("/styleTalent", summary="Style talent", description="Style talent")
 async def style_talent( talent_image: UploadFile = File(...), garment_image: UploadFile = File(...)):
     #Save the uploaded talent image tempoarily
     talent_image_bytes = await talent_image.read()
@@ -168,12 +217,13 @@ async def videoResult(request_id: str):
 
 #Auxiliary functions
 
-@app.post("/describeGarment")
+@app.post("/describeGarment", summary="Describe a garment or garments in the image",
+    description="Passes the garment image to openai to describe it to improve generation of base images")
 async def describe_garment(image: UploadFile = File(...)):
     image_bytes = await image.read()
     base64_image = base64.b64encode(image_bytes).decode("utf-8")
 
-    response = client.chat.completions.create(
+    response = openai_client.chat.completions.create(
         model="gpt-4o",
         messages=[
             {

requirements.txt

@@ -5,3 +5,4 @@ python-multipart
 fal-client
 requests
 openai
+redis