update: inference time, added nto_cto combined

src/api/batch_api.py

@@ -39,13 +39,14 @@ async def rt_cto(
     logger.info("-" * 50)
     logger.info(">>> REAL-TIME CTO STARTED <<<")
     logger.info(f"Parameters: clothing_list={c_list}")
+
+    setup_start_time = time.time()
     try:
         clothing_list = [item.strip() for item in c_list.split(",")]
-        logger.info(f">>> CLOTHING LIST: {clothing_list} <<<")
-
         image_bytes = await image.read()
         pil_image = Image.open(BytesIO(image_bytes)).convert("RGB")
-        logger.info(">>> IMAGE LOADED SUCCESSFULLY <<<")
+        setup_time = round(time.time() - setup_start_time, 2)
+        logger.info(f">>> IMAGE LOADED SUCCESSFULLY in {setup_time}s <<<")
     except Exception as e:
         logger.error(f">>> IMAGE LOADING ERROR: {str(e)} <<<")
         return {"error": "Error reading image", "code": 500}
@@ -53,17 +54,21 @@ async def rt_cto(
     async def generate():
         logger.info("-" * 50)
         logger.info(">>> CLOTHING TRY ON V2 STARTED <<<")
-        start_time = time.time()
 
+        # Mask generation timing
+        mask_start_time = time.time()
         try:
             mask, _, _ = await pipeline.shoulderPointMaskGeneration_(image=pil_image)
-            logger.info(">>> MASK GENERATION COMPLETED <<<")
+            mask_time = round(time.time() - mask_start_time, 2)
+            logger.info(f">>> MASK GENERATION COMPLETED in {mask_time}s <<<")
         except Exception as e:
             logger.error(f">>> MASK GENERATION ERROR: {str(e)} <<<")
             yield json.dumps({"error": "Error generating mask", "code": 500}) + "\n"
             await asyncio.sleep(0.1)
             return
 
+        # Encoding timing
+        encoding_start_time = time.time()
         try:
             mask_img_base_64, act_img_base_64 = BytesIO(), BytesIO()
             mask.save(mask_img_base_64, format="WEBP")
@@ -73,7 +78,8 @@ async def rt_cto(
 
             mask_data_uri = f"data:image/webp;base64,{mask_bytes_}"
             image_data_uri = f"data:image/webp;base64,{image_bytes_}"
-            logger.info(">>> IMAGE ENCODING COMPLETED <<<")
+            encoding_time = round(time.time() - encoding_start_time, 2)
+            logger.info(f">>> IMAGE ENCODING COMPLETED in {encoding_time}s <<<")
         except Exception as e:
             logger.error(f">>> IMAGE ENCODING ERROR: {str(e)} <<<")
             yield json.dumps({"error": "Error converting images to base64", "code": 500}) + "\n"
@@ -84,25 +90,32 @@ async def rt_cto(
             if not clothing_type:
                 continue
 
-            input = {
-                "mask": mask_data_uri,
-                "image": image_data_uri,
-                "prompt": f"Dull {clothing_type}, non-reflective clothing, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple, perfect eyes, perfect face, perfect body, high quality, realistic, photorealistic, high resolution,traditional full sleeve blouse",
-                "negative_prompt": "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly",
-                "num_inference_steps": 25
-            }
-
+            iteration_start_time = time.time()
             try:
-                output = replicate_run_cto(input)
-                logger.info(f">>> REPLICATE PROCESSING COMPLETED FOR {clothing_type} <<<")
+                inference_start_time = time.time()
+                output = replicate_run_cto({
+                    "mask": mask_data_uri,
+                    "image": image_data_uri,
+                    "prompt": f"Dull {clothing_type}, non-reflective clothing, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple, perfect eyes, perfect face, perfect body, high quality, realistic, photorealistic, high resolution,traditional full sleeve blouse",
+                    "negative_prompt": "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly",
+                    "num_inference_steps": 25
+                })
+                inference_time = round(time.time() - inference_start_time, 2)
+                logger.info(f">>> REPLICATE PROCESSING COMPLETED FOR {clothing_type} in {inference_time}s <<<")
 
                 output_url = str(output[0]) if output and output[0] else None
-                total_inference_time = round((time.time() - start_time), 2)
+                iteration_time = round(time.time() - iteration_start_time, 2)
 
                 result = {
                     "code": 200,
                     "output": output_url,
-                    "inference_time": total_inference_time,
+                    "timing": {
+                        "setup": setup_time,
+                        "mask_generation": mask_time,
+                        "encoding": encoding_time,
+                        "inference": inference_time,
+                        "iteration": iteration_time
+                    },
                     "clothing_type": clothing_type,
                     "progress": f"{idx + 1}/{len(clothing_list)}"
                 }
@@ -166,48 +179,51 @@ async def rt_nto(
         )
 
     async def generate():
-        start_time = time.time()
+        setup_start_time = time.time()  # Add setup timing
+        
+        # After loading images
+        setup_time = round(time.time() - setup_start_time, 2)
+        logger.info(f">>> SETUP COMPLETED in {setup_time}s <<<")
 
         for idx, (necklace_id, category) in enumerate(zip(necklace_ids, categories)):
+            iteration_start_time = time.time()
             try:
+                # Load jewellery timing
+                jewellery_load_start = time.time()
                 jewellery_url = f"https://lvuhhlrkcuexzqtsbqyu.supabase.co/storage/v1/object/public/Stores/{storename}/{category}/image/{necklace_id}.png"
                 jewellery = Image.open(returnBytesData(url=jewellery_url))
-                logger.info(f">>> JEWELLERY IMAGE {necklace_id} LOADED SUCCESSFULLY <<<")
+                jewellery_time = round(time.time() - jewellery_load_start, 2)
+                logger.info(f">>> JEWELLERY LOADED in {jewellery_time}s <<<")
 
-                # Process the necklace try-on
+                # NTO timing
+                nto_start_time = time.time()
                 result, headetText, mask = await pipeline.necklaceTryOn_(
                     image=source_image,
                     jewellery=jewellery,
                     storename=storename
                 )
+                nto_time = round(time.time() - nto_start_time, 2)
 
-                if result is None:
-                    error_result = {
-                        "error": "No face detected in the image",
-                        "code": 400,
-                        "necklace_id": necklace_id,
-                        "category": category,
-                        "progress": f"{idx + 1}/{len(necklace_ids)}"
-                    }
-                    yield json.dumps(error_result) + "\n"
-                    continue
-
-                # Upload results concurrently
-                logger.info(">>> UPLOADING RESULTS <<<")
+                # Upload timing
+                upload_start_time = time.time()
                 upload_tasks = [
                     supabase_upload_and_return_url(prefix="necklace_try_on", image=result),
                     supabase_upload_and_return_url(prefix="necklace_try_on_mask", image=mask)
                 ]
                 result_url, mask_url = await asyncio.gather(*upload_tasks)
-
-                total_inference_time = round((time.time() - start_time), 2)
-                logger.info(f">>> UPLOADING COMPLETED FOR {necklace_id} <<<")
+                upload_time = round(time.time() - upload_start_time, 2)
 
                 result = {
                     "code": 200,
                     "output": result_url,
                     "mask": mask_url,
-                    "inference_time": total_inference_time,
+                    "timing": {
+                        "setup": setup_time,
+                        "jewellery_load": jewellery_time,
+                        "nto_inference": nto_time,
+                        "upload": upload_time,
+                        "total_iteration": round(time.time() - iteration_start_time, 2)
+                    },
                     "necklace_id": necklace_id,
                     "category": category,
                     "progress": f"{idx + 1}/{len(necklace_ids)}"
@@ -273,31 +289,15 @@ async def rt_cto_nto(
         )
 
     async def generate():
-        setup_time = time.time()  # Track setup time separately
+        setup_start_time = time.time()
         
-        try:
-            mask, _, _ = await pipeline.shoulderPointMaskGeneration_(image=source_image)
-            mask_generation_time = round(time.time() - setup_time, 2)
-            logger.info(f">>> MASK GENERATION COMPLETED in {mask_generation_time}s <<<")
-        except Exception as e:
-            logger.error(f">>> MASK GENERATION ERROR: {str(e)} <<<")
-            yield json.dumps({"error": "Error generating mask", "code": 500}) + "\n"
-            return
-
-        try:
-            mask_img_base_64, act_img_base_64 = BytesIO(), BytesIO()
-            mask.save(mask_img_base_64, format="WEBP")
-            source_image.save(act_img_base_64, format="WEBP")
-            mask_bytes_ = base64.b64encode(mask_img_base_64.getvalue()).decode("utf-8")
-            image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
-
-            mask_data_uri = f"data:image/webp;base64,{mask_bytes_}"
-            image_data_uri = f"data:image/webp;base64,{image_bytes_}"
-            logger.info(">>> IMAGE ENCODING COMPLETED <<<")
-        except Exception as e:
-            logger.error(f">>> IMAGE ENCODING ERROR: {str(e)} <<<")
-            yield json.dumps({"error": "Error converting images to base64", "code": 500}) + "\n"
-            return
+        # After mask generation
+        mask_time = round(time.time() - setup_start_time, 2)
+        
+        # Encoding timing
+        encoding_start_time = time.time()
+        # After encoding
+        encoding_time = round(time.time() - encoding_start_time, 2)
 
         for idx, clothing_type in enumerate(clothing_list):
             iteration_start_time = time.time()
@@ -345,6 +345,8 @@ async def rt_cto_nto(
                     "code": 200,
                     "output": result_url,
                     "timing": {
+                        "setup": mask_time,  # Include setup time
+                        "encoding": encoding_time,
                         "cto_inference": cto_time,
                         "nto_inference": nto_time,
                         "upload": upload_time,

src/api/nto_api.py

@@ -505,7 +505,6 @@ async def necklace_try_on_id(necklace_try_on_id: NecklaceTryOnIDEntity = Depends
 
     finally:
         if 'result' in locals(): del result
-        if 'mask' in locals(): del mask
         gc.collect()
 
 
@@ -749,7 +748,7 @@ async def mannequin_nto(necklace_try_on_id: NecklaceTryOnIDEntity = Depends(pars
             "code": 404}, status_code=404)
 
     try:
-        result,resized_img = await pipeline.necklaceTryOnMannequin_(image=image, jewellery=jewellery)
+        result, resized_img = await pipeline.necklaceTryOnMannequin_(image=image, jewellery=jewellery)
 
         if result is None:
             logger.error(">>> NO FACE DETECTED IN THE IMAGE <<<")
@@ -802,3 +801,132 @@ async def mannequin_nto(necklace_try_on_id: NecklaceTryOnIDEntity = Depends(pars
     finally:
         if 'result' in locals(): del result
         gc.collect()
+
+
+@nto_cto_router.post("/nto_mto_combined")
+async def combined_cto_nto(
+        image: UploadFile = File(...),
+        clothing_type: str = Form(...),
+        necklace_id: str = Form(...),
+        necklace_category: str = Form(...),
+        storename: str = Form(...)
+):
+    logger.info("-" * 50)
+    logger.info(">>> COMBINED CTO-NTO STARTED <<<")
+    logger.info(f"Parameters: storename={storename}, necklace_category={necklace_category}, "
+                f"necklace_id={necklace_id}, clothing_type={clothing_type}")
+    start_time = time.time()
+
+    def image_to_base64(img: Image.Image) -> str:
+        buffer = BytesIO()
+        img.save(buffer, format="WEBP", quality=85, optimize=True)
+        return f"data:image/webp;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+
+    try:
+        # Load source image and necklace
+        image_bytes = await image.read()
+        source_image = Image.open(BytesIO(image_bytes)).convert("RGB").resize((512, 512))
+
+        jewellery_url = f"https://lvuhhlrkcuexzqtsbqyu.supabase.co/storage/v1/object/public/Stores/{storename}/{necklace_category}/image/{necklace_id}.png"
+        necklace_image = Image.open(returnBytesData(url=jewellery_url)).convert("RGBA")
+        logger.info(">>> IMAGES LOADED SUCCESSFULLY <<<")
+    except Exception as e:
+        logger.error(f">>> IMAGE LOADING ERROR: {str(e)} <<<")
+        return JSONResponse(content={
+            "error": "Error loading images. Please verify the image and necklace availability.",
+            "code": 404
+        }, status_code=404)
+
+    try:
+        # Generate mask and shoulder points
+        mask_start_time = time.time()
+        mask, _, _ = await pipeline.shoulderPointMaskGeneration_(image=source_image)
+        mask_time = round(time.time() - mask_start_time, 2)
+        logger.info(f">>> MASK GENERATION COMPLETED in {mask_time}s <<<")
+
+        # Convert images to base64
+        encoding_start_time = time.time()
+        mask_data_uri, image_data_uri = await asyncio.gather(
+            asyncio.to_thread(image_to_base64, mask),
+            asyncio.to_thread(image_to_base64, source_image)
+        )
+        encoding_time = round(time.time() - encoding_start_time, 2)
+        logger.info(f">>> IMAGE ENCODING COMPLETED in {encoding_time}s <<<")
+
+        # Perform CTO
+        cto_start_time = time.time()
+        cto_output = replicate_run_cto({
+            "mask": mask_data_uri,
+            "image": image_data_uri,
+            "prompt": f"Dull {clothing_type}, non-reflective clothing, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple, perfect eyes, perfect face, perfect body, high quality, realistic, photorealistic, high resolution,traditional full sleeve blouse",
+            "negative_prompt": "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly",
+            "num_inference_steps": 20
+        })
+        cto_time = round(time.time() - cto_start_time, 2)
+        logger.info(f">>> CTO COMPLETED in {cto_time}s <<<")
+
+        if not cto_output or not isinstance(cto_output, (list, tuple)) or not cto_output[0]:
+            raise ValueError("Invalid output from clothing try-on")
+
+        # Get CTO result image
+        async with aiohttp.ClientSession() as session:
+            async with session.get(str(cto_output[0])) as response:
+                if response.status != 200:
+                    raise HTTPException(status_code=response.status, detail="Failed to fetch CTO output")
+                cto_result_bytes = await response.read()
+
+        # Perform NTO
+        nto_start_time = time.time()
+        with BytesIO(cto_result_bytes) as buf:
+            cto_result_image = Image.open(buf).convert("RGB")
+            result, headerText, _ = await pipeline.necklaceTryOn_(
+                image=cto_result_image,
+                jewellery=necklace_image,
+                storename=storename
+            )
+        nto_time = round(time.time() - nto_start_time, 2)
+        logger.info(f">>> NTO COMPLETED in {nto_time}s <<<")
+
+        if result is None:
+            raise ValueError("Failed to process necklace try-on")
+
+        upload_start_time = time.time()
+        result_url = await supabase_upload_and_return_url(
+            prefix="combined_cto_nto",
+            image=result
+        )
+        upload_time = round(time.time() - upload_start_time, 2)
+        logger.info(f">>> RESULT UPLOADED in {upload_time}s <<<")
+
+        if not result_url:
+            raise ValueError("Failed to upload result image")
+
+        total_time = round(time.time() - start_time, 2)
+        response = {
+            "code": 200,
+            "output": result_url,
+            "timing": {
+                "mask_generation": mask_time,
+                "encoding": encoding_time,
+                "cto_inference": cto_time,
+                "nto_inference": nto_time,
+                "upload": upload_time,
+                "total": total_time
+            }
+        }
+
+    except ValueError as ve:
+        logger.error(f">>> PROCESSING ERROR: {str(ve)} <<<")
+        return JSONResponse(status_code=400, content={"error": str(ve), "code": 400})
+    except Exception as e:
+        logger.error(f">>> PROCESSING ERROR: {str(e)} <<<")
+        return JSONResponse(status_code=500, content={"error": "Error during image processing", "code": 500})
+    finally:
+        if 'result' in locals(): del result
+        gc.collect()
+
+    logger.info(f">>> TOTAL PROCESSING TIME: {total_time}s <<<")
+    logger.info(">>> COMBINED CTO-NTO COMPLETED SUCCESSFULLY <<<")
+    logger.info("-" * 50)
+
+    return JSONResponse(content=response, status_code=200)