Update app.py

app.py

@@ -33,11 +33,11 @@ def init_pool():
     port_inc()
 
 def pipe_t2i():
-    PIPE = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token"))
+    PIPE = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to("cuda")
     return PIPE
     
 def pipe_i2i():
-    PIPE = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+    PIPE = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
     PIPE.unet = torch.compile(PIPE.unet, mode="reduce-overhead", fullgraph=True)
     return PIPE