Update app.py

app.py

@@ -13,9 +13,11 @@ from PIL import Image
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
 from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
 
+
+
 # Define the device
 dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 # Use the 'waffles' environment variable as the access token
 hf_token = os.getenv('waffles')
@@ -30,18 +32,11 @@ 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(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")
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-max_memory = get_balanced_memory(
-    model,
-    max_memory=None,
-    no_split_module_classes=["DecoderLayer", "Attention", "MLP", "LayerNorm", "Linear"],
-    dtype='float16',
-    low_zero=False,
-)
 
 device_map = infer_auto_device_map(
     model,
@@ -50,7 +45,7 @@ device_map = infer_auto_device_map(
     dtype='float16'
 )
 
-model = dispatch_model(model, device_map='balanced')
+model = dispatch_model(model, device_map='torch.cuda:0')
 
 class calculateDuration:
     def __init__(self, activity_name=""):