loving zero gpu as much as i love taxes

app.py

@@ -7,8 +7,8 @@ from transformers import AutoModelForImageSegmentation
 from diffusers import FluxFillPipeline
 from PIL import Image, ImageOps
 
-# from sam2.sam2_image_predictor import SAM2ImagePredictor
-# import numpy as np
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+import numpy as np
 from simple_lama_inpainting import SimpleLama
 from contextlib import contextmanager
 
@@ -135,36 +135,35 @@ def rmbg(image=None, url=None):
 
 
 def mask_generation(image=None, d=None):
-    return image
-    # # use bfloat16 for the entire notebook
-    # # torch.autocast("cuda", dtype=torch.bfloat16).__enter__()
-    # # # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
-    # # if torch.cuda.get_device_properties(0).major >= 8:
-    # #     torch.backends.cuda.matmul.allow_tf32 = True
-    # #     torch.backends.cudnn.allow_tf32 = True
-    # d = eval(d)  # convert this to dictionary
-    # with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-    #     predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
-    #     predictor.set_image(image)
-    #     input_point = np.array(d["input_points"])
-    #     input_label = np.array(d["input_labels"])
-    #     masks, scores, logits = predictor.predict(
-    #         point_coords=input_point,
-    #         point_labels=input_label,
-    #         multimask_output=True,
-    #     )
-    # sorted_ind = np.argsort(scores)[::-1]
-    # masks = masks[sorted_ind]
-    # scores = scores[sorted_ind]
-    # logits = logits[sorted_ind]
-
-    # out = []
-    # for i in range(len(masks)):
-    #     m = Image.fromarray(masks[i] * 255).convert("L")
-    #     comp = Image.composite(image, m, m)
-    #     out.append((comp, f"image {i}"))
-
-    # return out
+    # use bfloat16 for the entire notebook
+    # torch.autocast("cuda", dtype=torch.bfloat16).__enter__()
+    # # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
+    # if torch.cuda.get_device_properties(0).major >= 8:
+    #     torch.backends.cuda.matmul.allow_tf32 = True
+    #     torch.backends.cudnn.allow_tf32 = True
+    d = eval(d)  # convert this to dictionary
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
+        predictor.set_image(image)
+        input_point = np.array(d["input_points"])
+        input_label = np.array(d["input_labels"])
+        masks, scores, logits = predictor.predict(
+            point_coords=input_point,
+            point_labels=input_label,
+            multimask_output=True,
+        )
+    sorted_ind = np.argsort(scores)[::-1]
+    masks = masks[sorted_ind]
+    scores = scores[sorted_ind]
+    logits = logits[sorted_ind]
+
+    out = []
+    for i in range(len(masks)):
+        m = Image.fromarray(masks[i] * 255).convert("L")
+        comp = Image.composite(image, m, m)
+        out.append((comp, f"image {i}"))
+
+    return out
 
 
 def erase(image=None, mask=None):