mods to get ZeroGPU working

app.py

@@ -42,15 +42,16 @@ def infer_mask_from_init_img(img, mask_with='grey'):
     "note, this works whether image is normalized on 0..1 or -1..1, but not 0..255"
     assert mask_with in ['blue','white','grey']
     "given an image with mask areas marked, extract the mask itself"
+    print("\n in infer_mask_from_init_img: ")
     if not torch.is_tensor(img):
         img = ToTensor()(img)
-    print("img.shape: ", img.shape)
+    print("    img.shape: ", img.shape)
     # shape of mask should be img shape without the channel dimension
     if len(img.shape) == 3:
         mask = torch.zeros(img.shape[-2:])
     elif len(img.shape) == 2:
         mask = torch.zeros(img.shape)
-    print("mask.shape: ", mask.shape)
+    print("    mask.shape: ", mask.shape)
     if mask_with == 'white':
         mask[ (img[0,:,:]==1) & (img[1,:,:]==1) & (img[2,:,:]==1)] = 1
     elif mask_with == 'blue':

sample.py

@@ -522,7 +522,7 @@ def get_init_image_and_mask(args, device):
     init_mask = init_mask.unsqueeze(0).unsqueeze(1).repeat(args.batch_size,3,1,1).float()
     return init_image.to(device), init_mask.to(device)
 
-@spaces.GPU
+#@spaces.GPU # generates an error
 def main():
     global init_image, init_mask
     p = argparse.ArgumentParser(description=__doc__,