Update app.py

app.py

@@ -70,6 +70,7 @@ current_date1 = get_current_date1()
 # Set environment variables for CUDA
 os.environ['PYTORCH_USE_CUDA_DSA'] = '1'
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 
 
 hf_token = os.getenv("HF_TOKEN")
@@ -711,6 +712,13 @@ def generate_map(location_names):
 #     image_3 = generate_image(hardcoded_prompt_3)
 #     return image_1, image_2, image_3
 
+
+
+
+# Clear any cached memory
+torch.cuda.empty_cache()
+
+
 import gradio as gr
 import torch
 from diffusers import FluxPipeline
@@ -720,8 +728,14 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 # Function to initialize Flux bot model
 def initialize_flux_bot():
-    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-    pipe.to(device)  # Move the model to the correct device (GPU/CPU)
+    try:
+        torch.cuda.empty_cache()  # Clear GPU memory cache
+        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)  # Use FP16
+        pipe.to(device)  # Move the model to the correct device (GPU/CPU)
+    except torch.cuda.OutOfMemoryError:
+        print("CUDA out of memory, switching to CPU.")
+        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float32)  # Use FP32 for CPU
+        pipe.to("cpu")
     return pipe
 
 # Function to generate image using Flux bot on the specified device
@@ -730,8 +744,8 @@ def generate_image_flux(prompt):
     image = pipe(
         prompt,
         guidance_scale=0.0,
-        num_inference_steps=4,
-        max_sequence_length=256,
+        num_inference_steps=2,
+        max_sequence_length=128,
         generator=torch.Generator(device).manual_seed(0)
     ).images[0]
     return image