fix: inference time

src/api/batch_api.py

@@ -273,11 +273,12 @@ async def rt_cto_nto(
         )
 
     async def generate():
-        start_time = time.time()
+        setup_time = time.time()  # Track setup time separately
         
         try:
             mask, _, _ = await pipeline.shoulderPointMaskGeneration_(image=source_image)
-            logger.info(">>> MASK GENERATION COMPLETED <<<")
+            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"
@@ -299,8 +300,10 @@ async def rt_cto_nto(
             return
 
         for idx, clothing_type in enumerate(clothing_list):
+            iteration_start_time = time.time()
             try:
                 # Perform CTO
+                cto_start_time = time.time()
                 cto_output = replicate_run_cto({
                     "mask": mask_data_uri,
                     "image": image_data_uri,
@@ -308,18 +311,17 @@ async def rt_cto_nto(
                     "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
                 })
-                
-                if not cto_output or not isinstance(cto_output, (list, tuple)) or not cto_output[0]:
-                    raise ValueError("Invalid output from clothing try-on")
+                cto_time = round(time.time() - cto_start_time, 2)
+                logger.info(f">>> CTO COMPLETED for {clothing_type} in {cto_time}s <<<")
 
-                # Get CTO result image
+                # Get CTO result and process NTO
+                nto_start_time = time.time()
                 async with aiohttp.ClientSession() as session:
                     async with session.get(str(cto_output[0])) as response:
                         if response.status != 200:
                             raise ValueError("Failed to fetch CTO output")
                         cto_result_bytes = await response.read()
 
-                # Process NTO
                 with BytesIO(cto_result_bytes) as buf:
                     cto_result_image = Image.open(buf).convert("RGB")
                     result, headerText, mask = await pipeline.necklaceTryOn_(
@@ -327,27 +329,33 @@ async def rt_cto_nto(
                         jewellery=jewellery,
                         storename=storename
                     )
-
-                if result is None:
-                    raise ValueError("Failed to process necklace try-on")
+                nto_time = round(time.time() - nto_start_time, 2)
+                logger.info(f">>> NTO COMPLETED for {clothing_type} in {nto_time}s <<<")
 
                 # Upload result
-                result_url = await supabase_upload_and_return_url(prefix="clothing_necklace_try_on", image=result)
-                if not result_url:
-                    raise ValueError("Failed to upload result image")
+                upload_start_time = time.time()
+                result_url = await supabase_upload_and_return_url(
+                    prefix="clothing_necklace_try_on", 
+                    image=result
+                )
+                upload_time = round(time.time() - upload_start_time, 2)
 
-                # Stream result
+                # Stream result with detailed timing
                 output_result = {
                     "code": 200,
                     "output": result_url,
-                    "inference_time": round((time.time() - start_time), 2),
+                    "timing": {
+                        "cto_inference": cto_time,
+                        "nto_inference": nto_time,
+                        "upload": upload_time,
+                        "total_iteration": round(time.time() - iteration_start_time, 2)
+                    },
                     "clothing_type": clothing_type,
                     "progress": f"{idx + 1}/{len(clothing_list)}"
                 }
                 yield json.dumps(output_result) + "\n"
                 await asyncio.sleep(0.1)
 
-                # Clean up
                 del result
                 gc.collect()