Update app.py

app.py

@@ -109,15 +109,18 @@ def predict(radio, dict, word_mask, prompt=""):
             img.cuda()
             print ("yes, CUDA is available here !! ")
         
-        model = model.to(torch.device(device))
-        img = img.to(torch.device(device))
-        prompt = labels.to(torch.device(device))
         #------------------
         
         word_masks = [word_mask]
         with torch.no_grad(): 
             #torch.cuda.amp.autocast(): #
             preds = model(img.repeat(len(word_masks),1,1,1), word_masks)[0]
+        
+        #model = model.to(torch.device(device))
+        img = img.to(torch.device(device))
+        prompt = labels.to(torch.device(device))
+        #---------
+        
         init_image = dict['image'].convert('RGB').resize((imgRes, imgRes))
         filename = f"{uuid.uuid4()}.png"
         plt.imsave(filename,torch.sigmoid(preds[0][0]))
@@ -144,7 +147,7 @@ def predict(radio, dict, word_mask, prompt=""):
             preds = model(img.repeat(len(word_masks),1,1,1), word_masks)[0]
 
         
-        model = model.to(torch.device(device))
+        #model = model.to(torch.device(device))
         img = img.to(torch.device(device))
         prompt = labels.to(torch.device(device))