Update app.py

app.py

@@ -32,7 +32,7 @@ with open('loras.json', 'r') as f:
 
 # Initialize the base model with authentication and specify the device
 # Initialize the base model with authentication and specify the device
-pipe = DiffusionPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=dtype, token=hf_token, trust_remote_code=True, device_map='balanced').to("torch.cuda:0")
+pipe = DiffusionPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=dtype, token=hf_token, trust_remote_code=True, device_map='balanced').to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048