Update app.py

app.py

@@ -713,39 +713,75 @@ def generate_map(location_names):
 #     return image_1, image_2, image_3
 
 
+# import torch
+# from diffusers import DiffusionPipeline
+# import numpy as np
+# import random
+# import gradio as gr
+
+# # Constants for device and dtype
+# dtype = torch.bfloat16
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# # Clear CUDA memory
+# torch.cuda.empty_cache()
+
+# # Load the Flux pipeline model
+# pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+
+# MAX_SEED = np.iinfo(np.int32).max
+# MAX_IMAGE_SIZE = 2048
+
+# # Function for inference using the Flux pipeline
+# def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4):
+#     if randomize_seed:
+#         seed = random.randint(0, MAX_SEED)
+#     generator = torch.Generator(device).manual_seed(seed)
+    
+#     # Generate the image using the pipeline
+#     image = pipe(
+#         prompt=prompt,
+#         width=width,
+#         height=height,
+#         num_inference_steps=num_inference_steps,
+#         generator=generator,
+#         guidance_scale=0.0
+#     ).images[0]
+    
+#     return image, seed
+
+
+
 import torch
 from diffusers import DiffusionPipeline
-import numpy as np
-import random
-import gradio as gr
+import os
 
-# Constants for device and dtype
-dtype = torch.bfloat16
+# Set PYTORCH_CUDA_ALLOC_CONF to avoid memory fragmentation
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+
+# Clear CUDA cache before loading the model
+torch.cuda.empty_cache()
+
+# Use a smaller dtype (e.g., torch.float16)
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the Flux pipeline model
+# Load the model with a smaller precision
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
-
-# Function for inference using the Flux pipeline
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+# Reduce inference steps to save memory
+def generate_image_flux(prompt, seed=42, width=1024, height=1024, num_inference_steps=2):
     generator = torch.Generator(device).manual_seed(seed)
-    
-    # Generate the image using the pipeline
     image = pipe(
         prompt=prompt,
         width=width,
         height=height,
-        num_inference_steps=num_inference_steps,
+        num_inference_steps=num_inference_steps,  # Reduce the number of inference steps
         generator=generator,
         guidance_scale=0.0
     ).images[0]
-    
-    return image, seed
+    return image
+
 
 # Existing image generation function, updated to use the new Flux pipeline
 def update_images():