Update app.py

app.py

@@ -59,7 +59,7 @@ model_id_or_path = "CompVis/stable-diffusion-v1-4"
 pipe = StableDiffusionInpaintingPipeline.from_pretrained(
     model_id_or_path,
     revision="fp16", 
-    torch_dtype=torch.long, #float16
+    torch_dtype=torch.float16, #float16
     use_auth_token=auth_token
 )
 #self.register_buffer('n_', ...)
@@ -87,7 +87,7 @@ def predict(radio, dict, word_mask, prompt=""):
     elif(radio == "type what to keep"):
         img = transform(dict["image"]).squeeze(0)
         word_masks = [word_mask]
-        with torch.no_grad():
+        with torch.cuda.amp.autocast(): #with torch.no_grad(): 
             preds = model(img.repeat(len(word_masks),1,1,1), word_masks)[0]
         init_image = dict['image'].convert('RGB').resize((imgRes, imgRes))
         filename = f"{uuid.uuid4()}.png"
@@ -101,7 +101,7 @@ def predict(radio, dict, word_mask, prompt=""):
     else:
         img = transform(dict["image"]).unsqueeze(0)
         word_masks = [word_mask]
-        with torch.no_grad():
+        with torch.cuda.amp.autocast(): #with torch.no_grad():
             preds = model(img.repeat(len(word_masks),1,1,1), word_masks)[0]
         init_image = dict['image'].convert('RGB').resize((imgRes, imgRes))
         filename = f"{uuid.uuid4()}.png"