Update main.py

main.py

@@ -88,18 +88,19 @@ 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'):
+def load_image_from_base64(base64_str: str, mode='RGB'):
     """Decodes a Base64 string and opens it as a PIL Image."""
-    if "," not in s:
-        return None
+    if "," in base64_str:
+        base64_str = base64_str.split(",")[1]
     try:
-        img_data = base64.b64decode(s.split(",")[1])
+        img_data = base64.b64decode(base64_str)
         img = Image.open(io.BytesIO(img_data))
-        return img.convert(m)
+        return img.convert(mode)  # Convert to the specified mode (e.g., 'RGB', 'L', etc.)
     except Exception as e:
         print(f"Error loading image from base64: {e}")
         return None
 
+
 # === API ENDPOINTS (CORRECT AND GUARANTEED) ===
 
 @app.get("/")
@@ -123,6 +124,7 @@ async def api_generate(request: Request, inputs: ApiInput):
     # Load person and fabric images from base64
     person = load_image_from_base64(inputs.person_base64)
     fabric = load_image_from_base64(inputs.fabric_base64)
+    mask = load_image_from_base64(inputs.mask_base64, mode='L')
     
     if person is None or fabric is None:
         raise HTTPException(status_code=400, detail="Could not decode base64 images.")