Update app.py

app.py

@@ -3,6 +3,9 @@ import os
 import gradio as gr
 import numpy as np
 import json
+from accelerate import dispatch_model, infer_auto_device_map
+from accelerate.utils import get_balanced_memory
+from torch.cuda.amp import autocast
 import torch
 import random
 import time
@@ -27,11 +30,30 @@ 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).to(device)
+pipe = DiffusionPipeline.from_pretrained("sayakpaul/FLUX.1-merged", torch_dtype=dtype, token=hf_token, trust_remote_code=True, device_map='auto').to(device)
 
 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,
+    max_memory=max_memory,
+    no_split_module_classes=["DecoderLayer", "Attention", "MLP", "LayerNorm", "Linear"],
+    dtype='float16'
+)
+
+model = dispatch_model(model, device_map=device_map)
+
+device = ‘cuda:0’
+
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -48,7 +70,7 @@ class calculateDuration:
         else:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=200)
 def generate_images(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, num_images, progress):
     generator = torch.Generator(device=device).manual_seed(seed)
     images = []