Update app.py

app.py

@@ -5,152 +5,109 @@ from PIL import Image
 import base64
 import io
 import os
-
 from segment_anything import sam_model_registry, SamPredictor
 from diffusers import StableDiffusionXLInpaintPipeline
 
-
-# ------------------- Load model -------------------
+# ------------------- Load Models -------------------
 MODEL_PATH = "sam_vit_b_01ec64.pth"
 
 if not os.path.exists(MODEL_PATH):
     os.system(f"wget https://dl.fbaipublicfiles.com/segment_anything/{MODEL_PATH}")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("🧠 Loading SAM model...")
 sam = sam_model_registry["vit_b"](checkpoint=MODEL_PATH)
 sam.to(device=device)
 predictor = SamPredictor(sam)
+print("✅ SAM loaded successfully!")
 
-# ------------------- Load SDXL Inpainting -------------------
-print("🪄 Loading Stable Diffusion XL Inpainting Model...")
-from diffusers import StableDiffusionXLInpaintPipeline
-
-sdxl_model_id = "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
+print("🎨 Loading SDXL Inpainting model...")
 pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
-    sdxl_model_id,
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    revision="fp16",
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    variant="fp16",
 )
 pipe = pipe.to(device)
-print("✅ SDXL Inpainting loaded successfully!")
+print("✅ SDXL loaded successfully!")
 
-# ------------------- Helper -------------------
-def decode_base64_image(image_base64: str):
-    """Decode base64 string menjadi PIL Image"""
-    image_data = base64.b64decode(image_base64)
-    image = Image.open(io.BytesIO(image_data)).convert("RGB")
-    return np.array(image)
+# ------------------- Helper Functions -------------------
+def np_to_pil(np_img):
+    return Image.fromarray((np_img * 255).astype(np.uint8)) if np_img.dtype == np.float32 else Image.fromarray(np_img)
 
-
-def encode_mask_to_base64(mask: np.ndarray):
-    """Encode mask numpy array menjadi base64 PNG"""
-    mask_image = Image.fromarray((mask * 255).astype(np.uint8))
+def pil_to_b64(image: Image.Image):
     buffered = io.BytesIO()
-    mask_image.save(buffered, format="PNG")
+    image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
+def decode_image(image):
+    if isinstance(image, str):
+        image_data = base64.b64decode(image)
+        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+    return np.array(image)
 
-# ------------------- Inference function -------------------
-def predict(image_base64, box=None, points=None, labels=None):
-    """
-    image_base64: string base64 dari gambar RGB
-    box: [x1, y1, x2, y2] (optional)
-    points: list of [x, y] (optional)
-    labels: list of 1/0 (optional)
-    """
-
+# ------------------- Main Pipeline -------------------
+def segment_and_inpaint(image, prompt, negative_prompt="", guidance_scale=7.5, steps=30, seed=42):
     try:
-        image_np = decode_base64_image(image_base64)
+        # Step 1: Segmentasi dengan SAM
+        image_np = np.array(image.convert("RGB"))
         predictor.set_image(image_np)
 
-        box_np = np.array(box) if box else None
-        points_np = np.array(points) if points else None
-        labels_np = np.array(labels) if labels else None
+        # Gunakan titik tengah gambar sebagai fokus sementara
+        h, w, _ = image_np.shape
+        points = np.array([[w // 2, h // 2]])
+        labels = np.array([1])
 
-        masks, scores, logits = predictor.predict(
-            point_coords=points_np,
-            point_labels=labels_np,
-            box=box_np,
+        masks, scores, _ = predictor.predict(
+            point_coords=points,
+            point_labels=labels,
             multimask_output=True
         )
 
         best_idx = np.argmax(scores)
         mask = masks[best_idx]
+        mask_pil = Image.fromarray((mask * 255).astype(np.uint8)).convert("L")
 
-        mask_base64 = encode_mask_to_base64(mask)
-        return {"mask_base64": mask_base64, "score": float(scores[best_idx])}
-
-    except Exception as e:
-        return {"error": str(e)}
-
-def inpaint_background(image_base64, mask_base64, prompt, negative_prompt="", guidance_scale=7.5, steps=30, seed=42):
-    """
-    Mengganti background berdasarkan prompt menggunakan model SDXL Inpainting.
-    image_base64: base64 dari gambar RGBA (foreground+alpha)
-    mask_base64: base64 dari mask (putih=area yang diganti)
-    """
-    try:
-        # Decode image
-        image = decode_base64_image(image_base64)
-        mask_data = base64.b64decode(mask_base64)
-        mask = Image.open(io.BytesIO(mask_data)).convert("L")
-
+        # Step 2: Inpainting Background
         generator = torch.manual_seed(int(seed))
         result = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
-            image=Image.fromarray(image),
-            mask_image=mask,
+            image=image,
+            mask_image=mask_pil,
             guidance_scale=float(guidance_scale),
             num_inference_steps=int(steps),
             generator=generator
         ).images[0]
 
-        buffered = io.BytesIO()
-        result.save(buffered, format="PNG")
-        result_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-        return {"result_base64": result_b64, "status": "✅ Success"}
+        return result
 
     except Exception as e:
         import traceback
-        return {"error": str(e), "traceback": traceback.format_exc()}
-
+        print(traceback.format_exc())
+        return f"❌ Error: {str(e)}"
 
 # ------------------- Gradio Interface -------------------
-demo = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Textbox(label="Image (Base64)", lines=5, placeholder="Base64 encoded image"),
-        gr.Textbox(label="Box [x1, y1, x2, y2]", placeholder="[100, 100, 300, 300]"),
-        gr.Textbox(label="Points (optional)", placeholder="[[120,150],[130,160]]"),
-        gr.Textbox(label="Labels (optional)", placeholder="[1,0]"),
-    ],
-    outputs="json",
-    title="SAM Segmentation API",
-    description="API untuk segmentasi gambar menggunakan Segment Anything Model (SAM)."
-)
-
-demo2 = gr.Interface(
-    fn=inpaint_background,
-    inputs=[
-        gr.Textbox(label="Image (Base64)", lines=3),
-        gr.Textbox(label="Mask (Base64)", lines=3),
-        gr.Textbox(label="Prompt", placeholder="beautiful beach background"),
-        gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry"),
-        gr.Slider(1, 20, value=7.5, step=0.5, label="Guidance Scale"),
-        gr.Slider(10, 50, value=30, step=5, label="Steps"),
-        gr.Number(value=42, label="Seed"),
-    ],
-    outputs="json",
-    title="SDXL Background Inpainting API",
-    description="API untuk mengganti background menggunakan Stable Diffusion XL Inpainting."
-)
-
-
-if __name__ == "__main__":
-    app = gr.TabbedInterface(
-        [demo, demo2],
-        ["SAM Segmentation", "Background Inpainting"]
+with gr.Blocks(title="SAM + SDXL Background Changer") as app:
+    gr.Markdown("## 🎨 Background Changer using SAM + Stable Diffusion XL Inpainting")
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Upload Image", type="pil")
+            prompt = gr.Textbox(label="Prompt (background description)", placeholder="a beach at sunset")
+            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry")
+            guidance_scale = gr.Slider(1, 15, value=7.5, step=0.5, label="Guidance Scale")
+            steps = gr.Slider(10, 50, value=30, step=5, label="Inference Steps")
+            seed = gr.Number(value=42, label="Random Seed")
+            submit_btn = gr.Button("✨ Change Background")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Result", type="pil")
+
+    submit_btn.click(
+        fn=segment_and_inpaint,
+        inputs=[input_image, prompt, negative_prompt, guidance_scale, steps, seed],
+        outputs=[output_image]
     )
-    app.launch(server_name="0.0.0.0", server_port=7860)
+
+app.launch(server_name="0.0.0.0", server_port=7860)