Update app.py

app.py

@@ -33,7 +33,7 @@ JS = """function () {
 }"""
 
 if torch.cuda.is_available():
-    pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
+    pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16).to(0)
     pipe.load_lora_weights(default_lora, weight_name = default_weight_name) # default load lora
     pipe.fuse_lora(lora_scale=0.9)
 
@@ -95,7 +95,7 @@ def generate_image(
     nums:int=1,
     progress=gr.Progress(track_tqdm=True)):
 
-    pipe.to(device="cuda")
+    pipe.to("cuda")
 
     if seed == -1:
         seed = random.randint(0, MAX_SEED)