Update app.py

app.py

@@ -9,7 +9,7 @@ schnell_model = "black-forest-labs/FLUX.1-schnell"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe_dev = DiffusionPipeline.from_pretrained(dev_model, torch_dtype=torch.bfloat16).to(device)
+pipe_dev = DiffusionPipeline.from_pretrained(dev_model, torch_dtype=torch.bfloat16)
 pipe_schnell = DiffusionPipeline.from_pretrained(
     schnell_model,
     text_encoder=pipe_dev.text_encoder,
@@ -17,10 +17,10 @@ pipe_schnell = DiffusionPipeline.from_pretrained(
     tokenizer=pipe_dev.tokenizer,
     tokenizer_2=pipe_dev.tokenizer_2,
     torch_dtype=torch.bfloat16
-).to(device)
-
+)
 @spaces.GPU
 def run_dev_hyper(prompt):
+    pipe_dev.to("cuda")
     repo_name = "ByteDance/Hyper-SD"
     ckpt_name = "Hyper-FLUX.1-dev-8steps-lora.safetensors"
     pipe_dev.load_lora_weights(hf_hub_download(repo_name, ckpt_name))
@@ -30,6 +30,7 @@ def run_dev_hyper(prompt):
 
 @spaces.GPU
 def run_dev_turbo(prompt):
+    pipe_dev.to("cuda")
     repo_name = "alimama-creative/FLUX.1-Turbo-Alpha"
     ckpt_name = "diffusion_pytorch_model.safetensors"
     pipe_dev.load_lora_weights(hf_hub_download(repo_name, ckpt_name))
@@ -39,6 +40,7 @@ def run_dev_turbo(prompt):
 
 @spaces.GPU
 def run_schnell(prompt):
+    pipe_schnell.to("cuda")
     image = pipe_schnell(prompt).images[0]
     return image