Create app.py

app.py

@@ -0,0 +1,54 @@
+import torch
+import numpy as np
+import cv2
+import gradio as gr
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+from tqdm import tqdm
+from diffusers import StableDiffusionInpaintPipeline, DDIMScheduler
+from PIL import Image
+
+# Load mô hình SAM
+sam = sam_model_registry["vit_h"](checkpoint="models/sam_vit_h_4b8939.pth")
+mask_generator = SamAutomaticMaskGenerator(sam)
+
+# Load mô hình Stable Diffusion
+scheduler = DDIMScheduler.from_pretrained("runwayml/stable-diffusion-inpainting", subfolder="scheduler")
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    scheduler=scheduler,
+    torch_dtype=torch.float32,
+    cache_dir="./models",
+    low_cpu_mem_usage=True
+).to("cuda" if torch.cuda.is_available() else "cpu")
+
+pipe.enable_attention_slicing()
+
+# Hàm xử lý ảnh
+def inpaint(image, mask, prompt):
+    # Convert to PIL
+    original_image = Image.fromarray(image).convert("RGB")
+    mask_image = Image.fromarray(mask).convert("L")
+
+    # Resize images to 512x512
+    original_image = original_image.resize((512, 512))
+    mask_image = mask_image.resize((512, 512))
+
+    # Inpainting AI
+    output = pipe(prompt=prompt, image=original_image, mask_image=mask_image, num_inference_steps=25).images[0]
+    return np.array(output)
+
+# UI với Gradio
+interface = gr.Interface(
+    fn=inpaint,
+    inputs=[
+        gr.Image(type="numpy", label="Upload Image"),
+        gr.Image(type="numpy", label="Upload Mask"),
+        gr.Textbox(label="Prompt (Describe what to add)")
+    ],
+    outputs=gr.Image(label="Generated Image"),
+    title="AI Furniture Inpainting",
+    description="Upload an image of a room and a mask where furniture should be added."
+)
+
+if __name__ == "__main__":
+    interface.launch()