Upload app.py

app.py

@@ -47,7 +47,7 @@ vae_processor = diffusers.image_processor.VaeImageProcessor(
 
 
 def mask_generation(image, processor, model, category):
-    inputs = processor(images=image, return_tensors="pt")
+    inputs = processor(images=image, return_tensors="pt").to("cuda")
     outputs = model(**inputs)
     logits = outputs.logits.cpu()
 
@@ -98,7 +98,7 @@ def process(
     image = image_data["background"]
     processor = SegformerImageProcessor.from_pretrained("sayeed99/segformer_b3_clothes")
     model = AutoModelForSemanticSegmentation.from_pretrained("sayeed99/segformer_b3_clothes")
-    model.to('cpu')
+    model.to("cuda")
 
     # preprocess image
     image = image.convert("RGB").resize((image_width, image_height))