Update app.py

app.py

@@ -10,22 +10,16 @@ from omegaconf import OmegaConf
 from huggingface_hub import snapshot_download
 
 REPO_DIR = snapshot_download(repo_id="Hang2991/PICS")
-# os.chdir(REPO_DIR)
+os.chdir(REPO_DIR)
 sys.path.insert(0, REPO_DIR)
 sys.path.insert(0, os.path.join(REPO_DIR, "dinov2"))
 
 from cldm.model import create_model, load_state_dict
 from cldm.ddim_hacked import DDIMSampler
 from datasets.data_utils import *
-# config = OmegaConf.load('configs/inference.yaml')
-config_path = os.path.join(REPO_DIR, 'configs/inference.yaml')
-config = OmegaConf.load(config_path)
-# pretrained_path = os.path.join(REPO_DIR, config.pretrained_model)
-
-full_config_file = os.path.join(REPO_DIR, config.config_file)
-full_pretrained_path = os.path.join(REPO_DIR, config.pretrained_model)
-model = create_model(full_config_file).cpu()
-model.load_state_dict(load_state_dict(full_pretrained_path, location='cpu'))
+config = OmegaConf.load('configs/inference.yaml')
+model = create_model(config.config_file).cpu()
+model.load_state_dict(load_state_dict(config.pretrained_model, location='cpu'))
 model.eval()
 
 def get_input(batch, k):
@@ -96,24 +90,17 @@ def process_composition(item, obj_thr):
 
 @spaces.GPU(duration=120)
 def pics_pairwise_inference(background, img_a, mask_a, img_b, mask_b):
-    def to_numpy(img):
-        if img is None: return None
-        # 如果 Gradio 传过来的是路径字符串，手动读取它
-        if isinstance(img, str):
-            return cv2.cvtColor(cv2.imread(img), cv2.COLOR_BGR2RGB)
-        return np.array(img)
-    
     device = "cuda"
     model.to(device)
     ddim_sampler = DDIMSampler(model)
-    back_image = to_numpy(background)
+    back_image = np.array(background)
     
     item_with_collage = {}
     objs = [(img_a, mask_a), (img_b, mask_b)]
     for j, (img, mask) in enumerate(objs):
         temp_patch = f"temp_obj_{j}.png"
-        cv2.imwrite(temp_patch, cv2.cvtColor(to_numpy(img), cv2.COLOR_RGB2BGR))
-        tar_mask = (to_numpy(mask)[:, :, 0] > 128).astype(np.uint8)
+        cv2.imwrite(temp_patch, cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+        tar_mask = (np.array(mask)[:, :, 0] > 128).astype(np.uint8)
         item_with_collage.update(process_pairs_multiple(tar_mask, back_image, temp_patch, counter=j))
 
     item_with_collage = process_composition(item_with_collage, obj_thr=2)
@@ -217,4 +204,4 @@ with gr.Blocks(title="PICS: Pairwise Spatial Compositing with Spatial Interactio
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(allowed_paths=[REPO_DIR])