Update app.py

app.py

@@ -75,7 +75,7 @@ device = 'cuda'
 
 base_model = "black-forest-labs/FLUX.1-dev"
 file_flux = hf_hub_download("marduk191/Flux.1_collection", "flux.1_dev_8x8_e4m3fn-marduk191.safetensors")
-pipe = FluxPipeline.from_single_file(file_flux, torch_dtype=torch.bfloat16, token=huggingface_token).to(device)
+pipe = FluxPipeline.from_single_file(file_flux, torch_dtype=torch.bfloat16, token=huggingface_token)
 
 # Load and fuse LoRA BEFORE quantizing
 print('Loading and fusing lora, please wait...')
@@ -147,6 +147,7 @@ ts_cutoff = 2
 
 @spaces.GPU
 def generate_flux_image(prompt, height, width, steps, scales, seed):
+    pipe.to(device)
     return pipe(
         prompt=prompt, 
         width=int(height),