Update main.py

main.py

@@ -89,47 +89,69 @@ def create_the_final_results(fabric: Image.Image, person: Image.Image, mask: Ima
     return results
 
 def load_image_from_base64(s: str, m: str = 'RGB'):
-    if "," not in s: return None
-    try: return Image.open(io.BytesIO(base64.b64decode(s.split(",")[1]))).convert(m)
-    except: return None
+    """Decodes a Base64 string and opens it as a PIL Image."""
+    if "," not in s:
+        return None
+    try:
+        img_data = base64.b64decode(s.split(",")[1])
+        img = Image.open(io.BytesIO(img_data))
+        return img.convert(m)
+    except Exception as e:
+        print(f"Error loading image from base64: {e}")
+        return None
 
 # === API ENDPOINTS (CORRECT AND GUARANTEED) ===
 
 @app.get("/")
-def root(): return {"status": "API server is running. Model will load on first call."}
-class ApiInput(BaseModel): person_base64: str; fabric_base64: str; mask_base64: Optional[str] = None
+def root(): 
+    return {"status": "API server is running. Model will load on first call."}
+
+class ApiInput(BaseModel):
+    person_base64: str
+    fabric_base64: str
+    mask_base64: Optional[str] = None
 
 @app.post("/generate")
 async def api_generate(request: Request, inputs: ApiInput):
-    print( "api_generate processing.." );
+    # Ensure the model is loaded
     load_model()
+    
     API_KEY = os.environ.get("API_KEY")
-    if request.headers.get("x-api-key") != API_KEY: raise HTTPException(status_code=401, detail="Unauthorized")
-    person = load_image_from_base64(inputs.person_base64); fabric = load_image_from_base64(inputs.fabric_base64)
-    if person is None or fabric is None: raise HTTPException(status_code=400, detail="Could not decode base64.")
+    if request.headers.get("x-api-key") != API_KEY:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+    
+    # Load person and fabric images from base64
+    person = load_image_from_base64(inputs.person_base64)
+    fabric = load_image_from_base64(inputs.fabric_base64)
+    
+    if person is None or fabric is None:
+        raise HTTPException(status_code=400, detail="Could not decode base64 images.")
     
-    # Process at high resolution, just like the Colab notebook.
+    # Resize person image to a standard size
     TARGET_SIZE = (1024, 1024)
     person_resized = person.resize(TARGET_SIZE, Image.Resampling.LANCZOS)
     
+    # Handle mask image if provided
     if inputs.mask_base64:
-        print( "api_generate if start processing.." );
         mask = load_image_from_base64(inputs.mask_base64, mode='L')
-        if mask is None: raise HTTPException(status_code=400, detail="Could not decode mask base64.")
-        # mask = mask.resize(TARGET_SIZE, Image.Resampling.LANCZOS)
-        print( "api_generate if end processing.." );
-    else: 
+        if mask is None:
+            raise HTTPException(status_code=400, detail="Could not decode mask base64.")
+        mask = mask.resize(TARGET_SIZE, Image.Resampling.LANCZOS)
+    else:
+        # If no mask is provided, generate one
         mask = generate_ultimate_mask(person_resized)
-        print( "api_generate else processing.." );
-        
+    
+    # Process and create the final results
     final_results = create_the_final_results(fabric, person_resized, mask)
     
+    # Convert image to base64 for the response
     def to_base64(img):
-        print( "to_base64 processing.." );
         img_display = img.resize((512, 512), Image.Resampling.LANCZOS)
-        buf = io.BytesIO(); img_display.save(buf, format="PNG"); 
+        buf = io.BytesIO()
+        img_display.save(buf, format="PNG")
         return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
         
+    # Prepare the response data with images converted to base64
     response_data = {
         'ultimate_image': to_base64(final_results['ultimate_image']),
         'fine_weave_image': to_base64(final_results['fine_weave_image']),
@@ -137,5 +159,5 @@ async def api_generate(request: Request, inputs: ApiInput):
         'creative_variation_image': to_base64(final_results['creative_variation_image']),
         'mask_image': to_base64(mask)
     }
-    print( "All Processing End" );
+    
     return response_data