Update app.py

app.py

@@ -1,21 +1,23 @@
-from fastapi import FastAPI, File, UploadFile, Form, BackgroundTasks
-from fastapi.responses import FileResponse, JSONResponse
+# app.py - AI Background Changer API
+
+from fastapi import FastAPI, File, UploadFile, Form, BackgroundTasks, HTTPException
+from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+from typing import Optional, List
 import numpy as np
 import os
 import uuid
 import cv2
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw
+import io
 import json
-from typing import List, Optional
 from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
-from huggingface_hub import hf_hub_download
 from diffusers import StableDiffusionXLInpaintPipeline
 
 # ======================================================
-# FASTAPI SETUP
+# CONFIGURATION
 # ======================================================
 app = FastAPI(title="AI Background Changer API")
 
@@ -31,82 +33,105 @@ app.add_middleware(
 # DIRECTORIES
 # ======================================================
 UPLOAD_DIR = "uploads"
-PROCESSED_DIR = "processed"
-MODELS_DIR = "models"
+RESULT_DIR = "results"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
-os.makedirs(PROCESSED_DIR, exist_ok=True)
-os.makedirs(MODELS_DIR, exist_ok=True)
+os.makedirs(RESULT_DIR, exist_ok=True)
 
-SAM_CHECKPOINT = os.path.join(MODELS_DIR, "sam_vit_b_01ec64.pth")
+MODEL_DIR = "models"
+SAM_CHECKPOINT = os.path.join(MODEL_DIR, "sam_vit_b_01ec64.pth")
 
 # ======================================================
-# GLOBAL STATE
+# GLOBAL STATE (In-memory session storage)
 # ======================================================
-global_models = {
-    "sam_predictor": None,
-    "inpaint_pipeline": None,
-    "device": "cuda" if torch.cuda.is_available() else "cpu"
-}
-
 sessions = {}
 
-# ======================================================
-# DOWNLOAD SAM FROM HUGGINGFACE
-# ======================================================
-def download_sam_checkpoint():
-    if not os.path.exists(SAM_CHECKPOINT):
-        print("Downloading SAM checkpoint from HuggingFace...")
-        hf_hub_download(
-            repo_id="yurista/AI-Background-Maker",
-            filename="models/sam_vit_b_01ec64.pth",
-            local_dir="models",
-            local_dir_use_symlinks=False
-        )
-        print("SAM checkpoint downloaded!")
-    return SAM_CHECKPOINT
+class SessionData:
+    def __init__(self):
+        self.original_image = None
+        self.current_mask = None
+        self.rgba_image = None
+        self.image_set = False
+        self.box_points = []
+        self.positive_points = []
+        self.negative_points = []
 
 # ======================================================
-# LOAD MODELS
+# MODELS
 # ======================================================
-def load_sam_model():
-    download_sam_checkpoint()
+sam_predictor = None
+inpaint_pipeline = None
 
-    if global_models["sam_predictor"] is None:
-        print(f"Loading SAM model from {SAM_CHECKPOINT}...")
+def load_sam_model():
+    global sam_predictor
+    if sam_predictor is None:
+        if not os.path.exists(SAM_CHECKPOINT):
+            raise FileNotFoundError(f"SAM model not found at {SAM_CHECKPOINT}")
+        
         sam = sam_model_registry["vit_b"](checkpoint=SAM_CHECKPOINT)
-        sam.to(global_models["device"])
-        predictor = SamPredictor(sam)
-        global_models["sam_predictor"] = predictor
-        print("SAM model loaded!")
-    return global_models["sam_predictor"]
-
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        sam.to(device=device)
+        
+        sam_predictor = SamPredictor(sam)
+        print(f"✅ SAM loaded on {device}")
+    
+    return sam_predictor
 
 def load_inpaint_pipeline():
-    if global_models["inpaint_pipeline"] is None:
-        print("Loading SDXL Inpaint model...")
-        pipeline = StableDiffusionXLInpaintPipeline.from_pretrained(
+    global inpaint_pipeline
+    if inpaint_pipeline is None:
+        print("🎨 Loading Stable Diffusion XL Inpainting...")
+        inpaint_pipeline = StableDiffusionXLInpaintPipeline.from_pretrained(
             "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            torch_dtype=torch.float16,
+            variant="fp16",
         )
-        pipeline = pipeline.to(global_models["device"])
-        global_models["inpaint_pipeline"] = pipeline
-        print("SDXL Inpaint loaded!")
-    return global_models["inpaint_pipeline"]
+        
+        if torch.cuda.is_available():
+            inpaint_pipeline = inpaint_pipeline.to("cuda")
+        
+        print("✅ Inpainting model loaded")
+    
+    return inpaint_pipeline
+
+# ======================================================
+# PYDANTIC MODELS
+# ======================================================
+class ClickRequest(BaseModel):
+    session_id: str
+    x: int
+    y: int
+    mode: str  # "box", "positive", "negative"
+
+class MaskRequest(BaseModel):
+    session_id: str
+
+class BackgroundRequest(BaseModel):
+    session_id: str
+    background_prompt: str
+    negative_prompt: Optional[str] = "low quality, blurry"
+    guidance_scale: Optional[float] = 7.0
+    num_steps: Optional[int] = 30
+    seed: Optional[int] = 42
 
 # ======================================================
-# HELPERS
+# HELPER FUNCTIONS
 # ======================================================
-def save_image(image_array, suffix=""):
-    unique_id = str(uuid.uuid4())
-    filename = f"{unique_id}{suffix}.png"
-    path = os.path.join(PROCESSED_DIR, filename)
+async def cleanup_files(paths: list):
+    for path in paths:
+        if os.path.exists(path):
+            os.remove(path)
+            print(f"🧹 File deleted: {path}")
 
-    if isinstance(image_array, np.ndarray):
-        cv2.imwrite(path, cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR))
-    else:
-        image_array.save(path)
+def get_session(session_id: str) -> SessionData:
+    if session_id not in sessions:
+        sessions[session_id] = SessionData()
+    return sessions[session_id]
 
-    return path
+def pil_to_bytes(img: Image.Image, format="PNG") -> bytes:
+    buf = io.BytesIO()
+    img.save(buf, format=format)
+    buf.seek(0)
+    return buf.getvalue()
 
 # ======================================================
 # ENDPOINTS
@@ -114,113 +139,448 @@ def save_image(image_array, suffix=""):
 @app.get("/")
 def home():
     return {
-        "message": "API Ready",
-        "endpoints": ["upload", "segment/box", "segment/points", "segment/auto-smart"]
+        "message": "✅ AI Background Changer API is running!",
+        "endpoints": {
+            "upload": "POST /upload - Upload image",
+            "click": "POST /click - Add prompt point",
+            "generate_mask": "POST /generate_mask - Generate segmentation mask",
+            "auto_segment": "POST /auto_segment - Auto segment all objects",
+            "auto_smart_select": "POST /auto_smart_select - Auto select foreground",
+            "reset": "POST /reset - Reset session prompts",
+            "generate_background": "POST /generate_background - Change background",
+            "get_preview": "GET /preview/{session_id} - Get current preview"
+        }
     }
 
 @app.post("/upload")
 async def upload_image(file: UploadFile = File(...)):
-    ext = os.path.splitext(file.filename)[1]
-    uid = str(uuid.uuid4())
-    path = os.path.join(UPLOAD_DIR, f"{uid}{ext}")
-
-    with open(path, "wb") as f:
-        f.write(await file.read())
-
-    img = cv2.imread(path)
-    if img is None:
-        return JSONResponse(status_code=400, content={"error": "Invalid image"})
-
-    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-    predictor = load_sam_model()
-    predictor.set_image(rgb)
-
-    sid = str(uuid.uuid4())
-    sessions[sid] = {
-        "image_path": path,
-        "mask": None,
-        "rgba_path": None,
-        "image_shape": rgb.shape
-    }
-
-    return {"session_id": sid, "image_shape": rgb.shape[:2]}
-
-@app.post("/segment/box")
-async def segment_box(
-    session_id: str = Form(...),
-    x1: int = Form(...),
-    y1: int = Form(...),
-    x2: int = Form(...),
-    y2: int = Form(...)
-):
-    if session_id not in sessions:
-        return JSONResponse(status_code=404, content={"error": "Session not found"})
-
-    sess = sessions[session_id]
-    img = cv2.imread(sess["image_path"])
-    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-    predictor = load_sam_model()
-    predictor.set_image(rgb)
-
-    box = np.array([min(x1,x2), min(y1,y2), max(x1,x2), max(y1,y2)])
-
-    masks, scores, _ = predictor.predict(box=box, multimask_output=True)
-    idx = np.argmax(scores)
-    mask = masks[idx]
-
-    sess["mask"] = mask.tolist()
-
-    rgba = np.dstack((rgb, (mask * 255).astype(np.uint8)))
-    rgba_path = save_image(rgba, "_rgba")
-    sess["rgba_path"] = rgba_path
-
-    segmented_path = save_image(rgba, "_segmented")
-
-    return {
-        "session_id": session_id,
-        "score": float(scores[idx]),
-        "segmented_preview": f"/files/{os.path.basename(segmented_path)}"
-    }
-
-
-@app.post("/generate-background")
+    """Upload image and initialize SAM"""
+    try:
+        # Create session
+        session_id = str(uuid.uuid4())
+        session = get_session(session_id)
+        
+        # Read image
+        contents = await file.read()
+        nparr = np.frombuffer(contents, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        
+        if image is None:
+            raise HTTPException(status_code=400, detail="Invalid image file")
+        
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        
+        # Store in session
+        session.original_image = image.copy()
+        session.image_set = False
+        
+        # Load SAM and set image
+        predictor = load_sam_model()
+        predictor.set_image(image)
+        session.image_set = True
+        
+        print(f"�� Image uploaded: {image.shape}, session: {session_id}")
+        
+        return {
+            "session_id": session_id,
+            "width": int(image.shape[1]),
+            "height": int(image.shape[0]),
+            "message": "Image uploaded successfully"
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/click")
+async def add_click(request: ClickRequest):
+    """Add click point (box, positive, or negative)"""
+    try:
+        session = get_session(request.session_id)
+        
+        if not session.image_set:
+            raise HTTPException(status_code=400, detail="Upload image first")
+        
+        if request.mode == "box":
+            if len(session.box_points) < 2:
+                session.box_points.append([request.x, request.y])
+            else:
+                session.box_points = [[request.x, request.y]]
+        elif request.mode == "positive":
+            session.positive_points.append([request.x, request.y])
+        elif request.mode == "negative":
+            session.negative_points.append([request.x, request.y])
+        
+        # Create preview with points
+        img_pil = Image.fromarray(session.original_image.copy())
+        draw = ImageDraw.Draw(img_pil)
+        
+        # Draw box
+        if len(session.box_points) > 0:
+            if len(session.box_points) == 1:
+                x, y = session.box_points[0]
+                draw.ellipse([x-8, y-8, x+8, y+8], fill=(0, 150, 255), outline=(255, 255, 255), width=3)
+            else:
+                x1, y1 = session.box_points[0]
+                x2, y2 = session.box_points[1]
+                draw.rectangle([x1, y1, x2, y2], outline=(0, 150, 255), width=4)
+                draw.ellipse([x1-8, y1-8, x1+8, y1+8], fill=(0, 150, 255), outline=(255, 255, 255), width=3)
+                draw.ellipse([x2-8, y2-8, x2+8, y2+8], fill=(0, 150, 255), outline=(255, 255, 255), width=3)
+        
+        # Draw positive points
+        for px, py in session.positive_points:
+            draw.ellipse([px-10, py-10, px+10, py+10], fill=(0, 255, 0), outline=(255, 255, 255), width=3)
+        
+        # Draw negative points
+        for nx, ny in session.negative_points:
+            draw.ellipse([nx-10, ny-10, nx+10, ny+10], fill=(255, 0, 0), outline=(255, 255, 255), width=3)
+        
+        img_bytes = pil_to_bytes(img_pil)
+        
+        return StreamingResponse(io.BytesIO(img_bytes), media_type="image/png")
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/generate_mask")
+async def generate_mask(request: MaskRequest):
+    """Generate segmentation mask from prompts"""
+    try:
+        session = get_session(request.session_id)
+        
+        if not session.image_set:
+            raise HTTPException(status_code=400, detail="Upload image first")
+        
+        has_box = len(session.box_points) == 2
+        has_points = len(session.positive_points) > 0
+        
+        if not has_box and not has_points:
+            raise HTTPException(status_code=400, detail="Add prompts first (box or points)")
+        
+        predictor = load_sam_model()
+        
+        box = None
+        points = None
+        labels = None
+        
+        # Box prompt
+        if has_box:
+            x1, y1 = session.box_points[0]
+            x2, y2 = session.box_points[1]
+            box = np.array([min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)])
+        
+        # Point prompts
+        if has_points or session.negative_points:
+            point_list = []
+            label_list = []
+            
+            for px, py in session.positive_points:
+                point_list.append([px, py])
+                label_list.append(1)
+            
+            for nx, ny in session.negative_points:
+                point_list.append([nx, ny])
+                label_list.append(0)
+            
+            points = np.array(point_list)
+            labels = np.array(label_list)
+        
+        # Predict
+        masks, scores, logits = predictor.predict(
+            point_coords=points,
+            point_labels=labels,
+            box=box,
+            multimask_output=True
+        )
+        
+        best_idx = np.argmax(scores)
+        mask = masks[best_idx]
+        score = scores[best_idx]
+        
+        session.current_mask = mask
+        
+        # Create RGBA
+        image = session.original_image
+        rgba_image = np.dstack((image, (mask * 255).astype(np.uint8)))
+        session.rgba_image = rgba_image
+        
+        # Convert to PIL
+        segmented_pil = Image.fromarray(rgba_image)
+        
+        mask_visual = (mask * 255).astype(np.uint8)
+        mask_visual_rgb = np.stack([mask_visual] * 3, axis=-1)
+        mask_pil = Image.fromarray(mask_visual_rgb)
+        
+        # Statistics
+        mask_area = mask.sum()
+        total_area = image.shape[0] * image.shape[1]
+        area_percentage = (mask_area / total_area) * 100
+        
+        # Return segmented image
+        img_bytes = pil_to_bytes(segmented_pil)
+        
+        return StreamingResponse(
+            io.BytesIO(img_bytes), 
+            media_type="image/png",
+            headers={
+                "X-Score": str(float(score)),
+                "X-Area-Percentage": str(float(area_percentage)),
+                "X-Method": "box" if has_box else "points"
+            }
+        )
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/auto_segment")
+async def auto_segment_all(request: MaskRequest):
+    """Auto segment and show all objects"""
+    try:
+        session = get_session(request.session_id)
+        
+        if not session.image_set:
+            raise HTTPException(status_code=400, detail="Upload image first")
+        
+        predictor = load_sam_model()
+        sam_model = predictor.model
+        
+        mask_generator = SamAutomaticMaskGenerator(
+            model=sam_model,
+            points_per_side=32,
+            pred_iou_thresh=0.86,
+            stability_score_thresh=0.92,
+            min_mask_region_area=100
+        )
+        
+        masks = mask_generator.generate(session.original_image)
+        masks = sorted(masks, key=lambda x: x['area'], reverse=True)
+        
+        # Create visualization
+        image = session.original_image.copy()
+        h, w = image.shape[:2]
+        overlay = np.zeros((h, w, 3), dtype=np.uint8)
+        
+        colors = [
+            [255, 0, 0], [0, 255, 0], [0, 0, 255],
+            [255, 255, 0], [255, 0, 255], [0, 255, 255],
+            [128, 0, 0], [0, 128, 0], [0, 0, 128]
+        ]
+        
+        objects_info = []
+        for i, mask_data in enumerate(masks[:9]):
+            mask = mask_data['segmentation']
+            color = colors[i % len(colors)]
+            overlay[mask] = color
+            
+            area_pct = (mask_data['area'] / (h * w)) * 100
+            objects_info.append({
+                "index": i + 1,
+                "area_percentage": float(area_pct),
+                "color": color
+            })
+        
+        result = cv2.addWeighted(image, 0.6, overlay, 0.4, 0)
+        result_pil = Image.fromarray(result)
+        
+        img_bytes = pil_to_bytes(result_pil)
+        
+        return StreamingResponse(
+            io.BytesIO(img_bytes),
+            media_type="image/png",
+            headers={
+                "X-Objects-Found": str(len(masks)),
+                "X-Objects-Info": json.dumps(objects_info)
+            }
+        )
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/auto_smart_select")
+async def auto_smart_select(request: MaskRequest):
+    """Automatically select the best foreground object"""
+    try:
+        session = get_session(request.session_id)
+        
+        if not session.image_set:
+            raise HTTPException(status_code=400, detail="Upload image first")
+        
+        predictor = load_sam_model()
+        sam_model = predictor.model
+        
+        mask_generator = SamAutomaticMaskGenerator(
+            model=sam_model,
+            points_per_side=32,
+            pred_iou_thresh=0.88,
+            stability_score_thresh=0.93,
+            min_mask_region_area=500
+        )
+        
+        masks = mask_generator.generate(session.original_image)
+        
+        if len(masks) == 0:
+            raise HTTPException(status_code=400, detail="No objects detected")
+        
+        # Smart selection
+        image = session.original_image
+        h, w = image.shape[:2]
+        center_x, center_y = w // 2, h // 2
+        
+        def score_mask(mask_data):
+            mask = mask_data['segmentation']
+            area = mask_data['area']
+            stability = mask_data['stability_score']
+            
+            y_coords, x_coords = np.where(mask)
+            if len(x_coords) == 0:
+                return 0
+            
+            mask_center_x = x_coords.mean()
+            mask_center_y = y_coords.mean()
+            
+            dist_from_center = np.sqrt((mask_center_x - center_x)**2 + (mask_center_y - center_y)**2)
+            max_dist = np.sqrt(center_x**2 + center_y**2)
+            center_score = 1 - (dist_from_center / max_dist)
+            
+            area_ratio = area / (h * w)
+            if area_ratio > 0.8:
+                size_score = 0.1
+            elif area_ratio < 0.02:
+                size_score = 0.3
+            else:
+                size_score = min(area_ratio * 5, 1.0)
+            
+            total_score = (center_score * 0.4 + size_score * 0.4 + stability * 0.2)
+            return total_score
+        
+        scored_masks = [(score_mask(m), m) for m in masks]
+        scored_masks.sort(reverse=True, key=lambda x: x[0])
+        
+        best_score, best_mask_data = scored_masks[0]
+        best_mask = best_mask_data['segmentation']
+        
+        session.current_mask = best_mask
+        
+        rgba_image = np.dstack((image, (best_mask * 255).astype(np.uint8)))
+        session.rgba_image = rgba_image
+        
+        segmented_pil = Image.fromarray(rgba_image)
+        
+        mask_area = best_mask.sum()
+        area_percentage = (mask_area / (h * w)) * 100
+        
+        img_bytes = pil_to_bytes(segmented_pil)
+        
+        return StreamingResponse(
+            io.BytesIO(img_bytes),
+            media_type="image/png",
+            headers={
+                "X-Selection-Score": str(float(best_score)),
+                "X-Stability": str(float(best_mask_data['stability_score'])),
+                "X-Area-Percentage": str(float(area_percentage)),
+                "X-Total-Objects": str(len(masks))
+            }
+        )
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/reset")
+async def reset_prompts(request: MaskRequest):
+    """Reset all prompts but keep the image"""
+    try:
+        session = get_session(request.session_id)
+        
+        session.box_points = []
+        session.positive_points = []
+        session.negative_points = []
+        session.current_mask = None
+        
+        return {"message": "Session reset successfully"}
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/generate_background")
 async def generate_background(
-    session_id: str = Form(...),
-    prompt: str = Form("futuristic city"),
-    negative_prompt: str = Form("blurry"),
-    guidance_scale: float = Form(7.0),
-    num_steps: int = Form(30)
+    background_tasks: BackgroundTasks,
+    request: BackgroundRequest
 ):
-    if session_id not in sessions:
-        return JSONResponse(status_code=404, content={"error": "Session not found"})
-
-    sess = sessions[session_id]
-    if sess["rgba_path"] is None:
-        return JSONResponse(status_code=400, content={"error": "Mask not generated"})
-
-    rgba = Image.open(sess["rgba_path"])
-    rgb = rgba.convert("RGB")
-    alpha = np.array(rgba.split()[-1])
-    mask = Image.fromarray(np.where(alpha < 128, 255, 0).astype(np.uint8))
-
-    pipe = load_inpaint_pipeline()
-
-    result = pipe(
-        image=rgb,
-        mask_image=mask,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=float(guidance_scale),
-        num_inference_steps=int(num_steps)
-    ).images[0]
-
-    out_path = save_image(result, "_bg")
-    return FileResponse(out_path, media_type="image/png")
-
-@app.get("/status/{session_id}")
-def status(session_id: str):
-    return sessions.get(session_id, {"error": "not found"})
-
-app.mount("/files", StaticFiles(directory=PROCESSED_DIR), name="files")
+    """Generate new background using Stable Diffusion XL Inpainting"""
+    try:
+        session = get_session(request.session_id)
+        
+        if session.rgba_image is None:
+            raise HTTPException(status_code=400, detail="Generate mask first")
+        
+        print(f"🎨 Generating background: {request.background_prompt}")
+        
+        # Load inpainting model
+        pipeline = load_inpaint_pipeline()
+        
+        # Prepare images
+        rgba_image = Image.fromarray(session.rgba_image)
+        alpha = np.array(rgba_image.split()[-1])
+        mask = np.where(alpha < 128, 255, 0).astype(np.uint8)
+        mask = Image.fromarray(mask).convert("L")
+        
+        rgb_image = rgba_image.convert("RGB")
+        
+        # Generate
+        generator = torch.manual_seed(int(request.seed))
+        result = pipeline(
+            image=rgb_image,
+            mask_image=mask,
+            prompt=request.background_prompt,
+            negative_prompt=request.negative_prompt,
+            guidance_scale=request.guidance_scale,
+            generator=generator,
+            num_inference_steps=int(request.num_steps),
+            width=rgb_image.width,
+            height=rgb_image.height
+        ).images[0]
+        
+        # Save result
+        result_id = str(uuid.uuid4())
+        result_path = os.path.join(RESULT_DIR, f"{result_id}.png")
+        result.save(result_path)
+        
+        # Schedule cleanup
+        background_tasks.add_task(cleanup_files, [result_path])
+        
+        return FileResponse(
+            result_path,
+            media_type="image/png",
+            headers={
+                "X-Prompt": request.background_prompt,
+                "X-Result-ID": result_id
+            }
+        )
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/preview/{session_id}")
+async def get_preview(session_id: str):
+    """Get current image preview"""
+    try:
+        session = get_session(session_id)
+        
+        if session.original_image is None:
+            raise HTTPException(status_code=400, detail="No image uploaded")
+        
+        img_pil = Image.fromarray(session.original_image)
+        img_bytes = pil_to_bytes(img_pil)
+        
+        return StreamingResponse(io.BytesIO(img_bytes), media_type="image/png")
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+def health_check():
+    return {
+        "status": "healthy",
+        "sam_loaded": sam_predictor is not None,
+        "inpaint_loaded": inpaint_pipeline is not None,
+        "device": "cuda" if torch.cuda.is_available() else "cpu"
+    }