update directed

app.py

@@ -1301,6 +1301,9 @@ def run_fn(
         model = load_alignedthreemodel()
     else:
         model = load_model(model_name)
+    
+    if directed:  # save qkv for directed, need more memory
+        model.enable_save_qkv()
         
     if "stable" in model_name.lower() and "diffusion" in model_name.lower():
         model.timestep = layer
@@ -1611,13 +1614,21 @@ def load_dataset_images(is_advanced, dataset_name, num_images=10,
         for i in valid_classes:
             idx = np.where(labels == i)[0]
             if is_random:
-                idx = np.random.RandomState(seed).choice(idx, chunk_size, replace=False)
+                if chunk_size < len(idx):
+                    idx = np.random.RandomState(seed).choice(idx, chunk_size, replace=False)
+                else:
+                    gr.Warning(f"Class {i} has less than {chunk_size} images.")
+                    idx = idx[:chunk_size]
             else:
                 idx = idx[:chunk_size]
             image_idx.extend(idx.tolist())
     if not is_filter:
         if is_random:
-            image_idx = np.random.RandomState(seed).choice(len(dataset), num_images, replace=False).tolist()
+            if num_images < len(dataset):
+                image_idx = np.random.RandomState(seed).choice(len(dataset), num_images, replace=False).tolist()
+            else:
+                gr.Warning(f"Dataset has less than {num_images} images.")
+                image_idx = list(range(num_images))
         else:
             image_idx = list(range(num_images))
     key = 'image' if 'image' in dataset[0] else list(dataset[0].keys())[0]
@@ -3303,21 +3314,18 @@ with demo:
                     # convert to PIL
                     mask = positive_mask.cpu().numpy()
                     mask = mask.astype(np.uint8) * 255
-                    mask = [Image.fromarray(mask[i]) for i in range(len(mask))]
                     
                     import cv2
-                    def get_bboxes_and_clean_mask(pil_mask, min_area=500):
+                    def get_bboxes_and_clean_mask(mask, min_area=500):
                         """
                         Args:
-                        - pil_mask: A Pillow image of a binary mask with 255 for the object and 0 for the background.
+                        - mask: A numpy image of a binary mask with 255 for the object and 0 for the background.
                         - min_area: Minimum area for a connected component to be considered valid (default 500).
                         
                         Returns:
                         - bounding_boxes: List of bounding boxes for valid objects (x, y, width, height).
                         - cleaned_pil_mask: A Pillow image of the cleaned mask, with small components removed.
                         """
-                        # Convert the Pillow image to a NumPy array
-                        mask = np.array(pil_mask)
 
                         # Ensure the mask is binary (0 or 255)
                         mask = np.where(mask > 127, 255, 0).astype(np.uint8)

requirements.txt

@@ -21,4 +21,4 @@ timm==0.9.2
 open-clip-torch==2.20.0
 pytorch_lightning==1.9.4
 gradio_image_prompter==0.1.0
-ncut-pytorch>=1.4.2
+ncut-pytorch>=1.5.2