wip

app.py

@@ -7,6 +7,8 @@ from transformers import AutoModelForImageSegmentation
 
 model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
 torch.set_float32_matmul_precision(['high', 'highest'][0])
+if torch.cuda.is_available():
+    model = model.to('cuda')
 model.eval()
 
 def remove_background(input_image, holiday, message):
@@ -58,8 +60,8 @@ demo = gr.Interface(
         gr.Image(type="pil", label="Second Output"),
         gr.Image(type="pil", label="Third Output")
     ],
-    title="Background Removal Tool",
-    description="Upload an image to remove its background"
+    title="Holiday Card Generator",
+    description="Upload an image to generate a holiday card"
 )
 
 demo.launch()