Update app.py

app.py

@@ -30,7 +30,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='auto').to(device)
+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
@@ -50,7 +50,7 @@ device_map = infer_auto_device_map(
     dtype='float16'
 )
 
-model = dispatch_model(model, device_map=device_map)
+model = dispatch_model(model, device_map='balanced')
 
 class calculateDuration:
     def __init__(self, activity_name=""):