Update app.py

app.py

@@ -9,14 +9,11 @@ from huggingface_hub import hf_hub_download, login
 import os
 import time
 
-# 🛠 Kiểm tra thiết bị
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 🔑 Đăng nhập Hugging Face
 hf_token = os.getenv("HF_TOKEN")
 login(token=hf_token)
 
-# 🏗 Tải mô hình SAM từ Hugging Face
 model_path = hf_hub_download(
     repo_id="Vuvo11/segment_anything_model",
     filename="sam_vit_h_4b8939.pth",
@@ -25,65 +22,56 @@ model_path = hf_hub_download(
 sam = sam_model_registry["vit_h"](checkpoint=model_path)
 mask_generator = SamAutomaticMaskGenerator(sam)
 
-# 🏗 Tải mô hình Stable Diffusion từ Hugging Face
 scheduler = DDIMScheduler.from_pretrained("runwayml/stable-diffusion-inpainting", subfolder="scheduler")
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
     "runwayml/stable-diffusion-inpainting",
     scheduler=scheduler,
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # 🔥 FP16 cho GPU
     cache_dir="./models",
     low_cpu_mem_usage=True
 ).to(device)
 
+if torch.cuda.is_available():
+    pipe.unet = torch.compile(pipe.unet)  # 🔥 Tối ưu tốc độ nếu chạy trên GPU
+
 pipe.enable_attention_slicing()
 
-# 🖼 Hàm tự động tạo mask từ ảnh gốc
-def generate_mask(image, progress=gr.Progress()):
-    progress(0, "Generating mask...")
+def generate_mask(image):
     masks = mask_generator.generate(image)
     if len(masks) == 0:
         return np.zeros_like(image[:, :, 0])
 
     largest_mask = max(masks, key=lambda x: np.sum(x["segmentation"]))
-    progress(50, "Mask generated successfully!")
     return (largest_mask["segmentation"] * 255).astype(np.uint8)
 
-# 🖌 Hàm xử lý ảnh (có progress)
 def inpaint(image, prompt, progress=gr.Progress()):
     progress(0, "Processing image...")
 
-    mask = generate_mask(image, progress)  # 🏗 Tạo mask tự động
+    mask = generate_mask(image)  
     progress(30, "Generating inpainting...")
 
     original_image = Image.fromarray(image).convert("RGB")
     mask_image = Image.fromarray(mask).convert("L")
 
-    original_image = original_image.resize((512, 512))
-    mask_image = mask_image.resize((512, 512))
+    original_image = original_image.resize((384, 384))  # 🔥 Resize nhỏ hơn để xử lý nhanh hơn
+    mask_image = mask_image.resize((384, 384))
 
-    for step in tqdm(range(25), desc="Inpainting Progress"):
-        time.sleep(0.1)
+    output = pipe(prompt=prompt, image=original_image, mask_image=mask_image, num_inference_steps=15).images[0]  # 🔥 Giảm số bước suy luận
 
-    output = pipe(prompt=prompt, image=original_image, mask_image=mask_image, num_inference_steps=25).images[0]
     progress(100, "Completed!")
     return np.array(output)
 
-# 🌐 UI với Gradio (có nút Submit)
 with gr.Blocks() as interface:
-    gr.Markdown("## 🎨 AI Furniture Inpainting")
+    gr.Markdown("## 🎨 AI Furniture Inpainting (Optimized)")
     
     with gr.Row():
         image_input = gr.Image(type="numpy", label="Upload Image")
         prompt_input = gr.Textbox(label="Prompt (Describe what to add)")
     
-    submit = gr.Button("Submit")  # 👉 Nút Submit để chạy
+    submit = gr.Button("Submit")
     output_image = gr.Image(label="Generated Image")
 
-    submit.click(
-        fn=inpaint,
-        inputs=[image_input, prompt_input],
-        outputs=output_image
-    )
+    submit.click(fn=inpaint, inputs=[image_input, prompt_input], outputs=output_image)
 
 if __name__ == "__main__":
     interface.launch()