Update main.py

main.py

@@ -32,56 +32,62 @@ def load_model():
     AI_MODEL["predictor"] = SamPredictor(sam)
     print("✅ High-Quality AI Model is now loaded.")
 
-# === CORE PROCESSING FUNCTIONS (DIRECTLY FROM YOUR WORKING COLAB) ===
+# === CORE PROCESSING FUNCTIONS (DIRECTLY FROM YOUR PERFECTED COLAB) ===
 
 def generate_ultimate_mask(image: Image.Image):
-    print( "Generating Mask image . . . ." );
-    """Generates the high-quality mask FOR THE SUIT ONLY, ignoring the shirt and pin."""
+    """Generates the high-quality mask FOR THE SUIT ONLY."""
+    print("   - DEBUG: Entering generate_ultimate_mask function...")
     sam_predictor = AI_MODEL["predictor"]; np = AI_MODEL["numpy"]
     image_np = np.array(image.convert('RGB')); sam_predictor.set_image(image_np); h, w, _ = image_np.shape
     input_points = np.array([[w*0.30,h*0.50],[w*0.70,h*0.50],[w*0.50,h*0.40],[w*0.65,h*0.32]])
     input_labels = np.array([1, 1, 0, 0])
     masks, _, _ = sam_predictor.predict(point_coords=input_points, point_labels=input_labels, multimask_output=False)
+    print("   - DEBUG: Mask generation complete.")
     return Image.fromarray(masks[0]).convert('L').filter(ImageFilter.GaussianBlur(2))
 
-def create_the_final_result(fabric: Image.Image, person: Image.Image, mask: Image.Image):
-    """THE FINAL, GUARANTEED, PIXEL-PERFECT COMPOSITING FUNCTION."""
-    # 1. Tile the fabric (Base Layer with perfect color).
-    print( "# 1. Tile the fabric (Base Layer with perfect color)." );
-    base_size=int(person.width/4); sw=max(1,int(base_size*0.65)); fw,fh=fabric.size
-    sh=max(1,int(fh*(sw/fw))) if fw>0 else 0; s=fabric.resize((sw,sh),Image.Resampling.LANCZOS)
-    tiled_fabric=Image.new('RGB',person.size)
-    for i in range(0,person.width,sw):
-      for j in range(0,person.height,sh): tiled_fabric.paste(s,(i,j))
-        
-    # 2. Create the Form & Shading Layer.
-    print( " # 2. Create the Form & Shading Layer." );
-    form_map=ImageOps.autocontrast(ImageOps.grayscale(person),cutoff=2).convert('RGB')
-    shaped_fabric=ImageChops.soft_light(tiled_fabric,form_map)
+def create_pixel_perfect_results(fabric: Image.Image, person: Image.Image, mask: Image.Image):
+    """The definitive, multi-layer professional compositing workflow."""
+    print("   - DEBUG: Entering create_pixel_perfect_results function...")
+    results = {}
+    
+    print("   - DEBUG: Creating shadow and highlight maps...")
+    grayscale_person = ImageOps.grayscale(person)
+    shadow_map = ImageOps.autocontrast(grayscale_person, cutoff=(10, 99)).convert('RGB')
+    highlight_map = ImageOps.invert(ImageOps.autocontrast(grayscale_person, cutoff=(90, 99))).convert('RGB')
 
-    # 3. Create the Detail & Contrast Layer.
-    print( " # 3. Create the Detail & Contrast Layer. " );
-    grayscale_person_details=ImageOps.grayscale(person)
-    shadow_map=ImageOps.autocontrast(grayscale_person_details,cutoff=(5,95)).convert('RGB')
-    highlight_map=ImageOps.invert(ImageOps.autocontrast(grayscale_person_details,cutoff=(95,99))).convert('RGB')
+    scales = {"classic": 0.75, "fine": 0.4, "bold": 1.2}
     
-    # 4. Apply the Detail Layers with Opacity.
-    print( " # 4. Apply the Detail Layers with Opacity." );
-    shadowed_layer=ImageChops.multiply(shaped_fabric,shadow_map)
-    final_shadows=Image.blend(shaped_fabric,shadowed_layer,alpha=0.25)
-    highlighted_layer=ImageChops.screen(final_shadows,highlight_map)
-    final_lit=Image.blend(final_shadows,highlighted_layer,alpha=0.1)
+    for style, sf in scales.items():
+        print(f"   - DEBUG: Generating '{style}' image...")
+        base_size = int(person.width / 4); sw = max(1, int(base_size * sf)); fw, fh = fabric.size
+        sh = max(1, int(fh * (sw / fw))) if fw > 0 else 0
+        s = fabric.resize((sw, sh), Image.Resampling.LANCZOS); tiled_fabric = Image.new('RGB', person.size)
+        for i in range(0, person.width, sw):
+            for j in range(0, person.height, sh): tiled_fabric.paste(s, (i, j))
+        
+        shadowed_layer = ImageChops.multiply(tiled_fabric, shadow_map)
+        shadowed_fabric = Image.blend(tiled_fabric, shadowed_layer, alpha=0.65)
+        highlighted_layer = ImageChops.screen(shadowed_fabric, highlight_map)
+        lit_fabric = Image.blend(shadowed_fabric, highlighted_layer, alpha=0.35)
+        
+        final_image = person.copy(); final_image.paste(lit_fabric, (0, 0), mask=mask)
+        results[f"{style}_image"] = final_image
 
-    # 5. Composite the final result.
-    print( " # 5. Composite the final result." );
-    final_image=person.copy(); final_image.paste(final_lit,(0,0),mask=mask)
-    return final_image
+    print("   - DEBUG: Generating 'realistic' (creative) image...")
+    light_map_rgb = ImageOps.autocontrast(ImageOps.grayscale(person).convert('RGB'), cutoff=2)
+    results["realistic_image"] = ImageChops.soft_light(results["classic_image"], light_map_rgb)
+    print("   - DEBUG: All image generation complete.")
+    return results
 
 def load_image_from_base64(s: str, m: str = 'RGB'):
-    print( "Load image from base64" );
-    if "," not in s: return None
-    try: return Image.open(io.BytesIO(base64.b64decode(s.split(",")[1]))).convert(m)
-    except: return None
+    """Decodes Base64 to a PIL Image. Fast and reliable."""
+    if not s or "," not in s: return None
+    try:
+        img_data = base64.b64decode(s.split(",")[1])
+        return Image.open(io.BytesIO(img_data)).convert(m)
+    except Exception as e:
+        print(f"   - ❌ ERROR: Failed to decode base64 string. Error: {e}")
+        return None
 
 # === API ENDPOINTS (CORRECT AND GUARANTEED) ===
 
@@ -95,35 +101,41 @@ class ApiInput(BaseModel):
 
 @app.post("/generate")
 async def api_generate(request: Request, inputs: ApiInput):
-    print( "---------------------------------Starting---------------------------------" );
-    load_model()
+    print("\n--- Received API request for /generate ---")
+    load_model() # This is the magic. All heavy lifting happens here, only once.
+    
     API_KEY = os.environ.get("API_KEY")
-    if request.headers.get("x-api-key") != API_KEY: raise HTTPException(status_code=401, detail="Unauthorized")
+    if request.headers.get("x-api-key") != API_KEY:
+        print("   - ❌ ERROR: Unauthorized. Invalid API Key provided.")
+        raise HTTPException(status_code=401, detail="Unauthorized")
     
+    print("   - DEBUG: Loading 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.")
+    if person is None or fabric is None:
+        raise HTTPException(status_code=400, detail="Could not decode person or fabric base64 string.")
     
-    #person_resized = person.resize((1024, 1024), Image.Resampling.LANCZOS)
-    person_resized = person
+    print("   - DEBUG: Resizing person image...")
+    person_resized = person.resize((1024, 1024), Image.Resampling.LANCZOS)
     
     if inputs.mask_base64:
+        print("   - DEBUG: Loading provided mask from Base64...")
         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((1024, 1024), Image.Resampling.LANCZOS)
+        mask = mask.resize((1024, 1024), Image.Resampling.LANCZOS)
     else: 
+        print("   - DEBUG: No mask provided, generating one automatically...")
         mask = generate_ultimate_mask(person_resized)
         
-    final_result = create_the_final_result(fabric, person_resized, mask)
+    result_images = create_pixel_perfect_results(fabric, person_resized, mask)
     
     def to_base64(img):
-        # Resize for a consistent output size
-        # img = img.resize((512, 512), Image.Resampling.LANCZOS)
+        img = img.resize((512, 512), Image.Resampling.LANCZOS)
         buf = io.BytesIO(); img.save(buf, format="PNG"); return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
-        
-    response_data = {
-        "result_image": to_base64(final_result),
-        "mask_image": to_base64(mask)
-    }
-    print( "---------------------------------Completed---------------------------------" );
+    
+    print("   - DEBUG: Encoding final images to Base64...")
+    response_data = {key: to_base64(img) for key, img in result_images.items()}
+    response_data["mask_image"] = to_base64(mask)
+    
+    print("--- ✅ Request successful. Sending response. ---")
     return response_data