Update app.py

app.py

@@ -14,7 +14,10 @@ models = {}
 for model_name in model_list:
     try:
         print(f"\n \nLoading {model_name}...")
-        models[model_name] = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16).to("cuda")
+        pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16).to("cuda")
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe.enable_model_cpu_offload()
+        models[model_name] = pipe
     except Exception as e:
         print(f"Error loading model {model_name}: {e}")
 
@@ -39,6 +42,7 @@ def generate_images(
         outputs = []
 
         for _ in range(num_images):
+            pipe.enable_model_cpu_offload()
             output = pipe(
                 prompt,
                 negative_prompt=negative_prompt,
@@ -48,6 +52,8 @@ def generate_images(
                 width=width
             )["images"][0]
             outputs.append(output)
+            
+