edit sam model source

app.py

@@ -6,12 +6,17 @@ from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 from tqdm import tqdm
 from diffusers import StableDiffusionInpaintPipeline, DDIMScheduler
 from PIL import Image
+from huggingface_hub import hf_hub_download
 
-# Load mô hình SAM
-sam = sam_model_registry["vit_h"](checkpoint="models/sam_vit_h_4b8939.pth")
+# 🛠 Kiểm tra thiết bị
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 🏗 Tải mô hình SAM từ Hugging Face
+sam_checkpoint = hf_hub_download(repo_id="facebook/sam", filename="sam_vit_h_4b8939.pth")
+sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
 mask_generator = SamAutomaticMaskGenerator(sam)
 
-# Load mô hình Stable Diffusion
+# 🏗 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",
@@ -19,7 +24,7 @@ pipe = StableDiffusionInpaintPipeline.from_pretrained(
     torch_dtype=torch.float32,
     cache_dir="./models",
     low_cpu_mem_usage=True
-).to("cuda" if torch.cuda.is_available() else "cpu")
+).to(device)
 
 pipe.enable_attention_slicing()
 

requirements.txt

@@ -9,4 +9,5 @@ pillow
 tqdm
 diffusers
 transformers
-segment-anything
+segment-anything
+huggingface_hub