Update app.py

app.py

@@ -14,7 +14,7 @@ from src.flux.xflux_pipeline import XFluxSampler
 args = OmegaConf.load("inference_configs/inference.yaml")
 # is_schnell = args.model_name == "flux-schnell"
 # sampler = None
-# device = torch.device("cuda")
+device = torch.device("cuda")
 # dtype = torch.bfloat16
 # dit = load_flow_model2(args.model_name, device="cpu").to(device, dtype=dtype)
 # vae = load_ae(args.model_name, device="cpu").to(device, dtype=dtype)
@@ -91,6 +91,7 @@ def generate(image: Image.Image, edit_prompt: str):
     img = img.permute(2, 0, 1).unsqueeze(0).to(device, dtype=dtype)
 
     result = sampler(
+        device='cuda',
         prompt=edit_prompt,
         width=args.sample_width,
         height=args.sample_height,