Update app.py

app.py

@@ -13,6 +13,8 @@ import base64
 from io import BytesIO
 import difflib
 
+# Constants
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # Helper functions
 def image_to_base64(image: Image.Image):
@@ -90,6 +92,7 @@ def getSegments(yoloModel, img1):
 
 @spaces.GPU
 def getDescript(image_captioner, img1):
+    image_captioner.model.to('cuda')
     base64_img = image_to_base64(img1)
     caption = image_captioner(base64_img)[0]['generated_text']
     return caption
@@ -110,9 +113,10 @@ def rmGPT(caption, remove_class, change):
         return ' '.join(arstr)
 
 @spaces.GPU
-def ChangeOBJ(sdxl_m, img1, response, mask1):
+def ChangeOBJ(sdxl_model, img1, response, mask1):
+    sdxl_model.to('cuda')
     size = img1.size
-    image = sdxl_m(prompt=response, image=img1, mask_image=mask1).images[0]
+    image = sdxl_model(prompt=response, image=img1, mask_image=mask1).images[0]
     return image.resize((size[0], size[1]))
 
 # Load models initially