Update app.py

app.py

@@ -756,27 +756,36 @@ import torch
 from diffusers import DiffusionPipeline
 import os
 
-# Set PYTORCH_CUDA_ALLOC_CONF to avoid memory fragmentation
+# Set PYTORCH_CUDA_ALLOC_CONF to handle memory fragmentation
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 
-# Clear CUDA cache before loading the model
-torch.cuda.empty_cache()
+# Check GPU memory and fallback to CPU if necessary
+if torch.cuda.is_available():
+    free_memory = torch.cuda.memory_reserved(0) - torch.cuda.memory_allocated(0)
+    if free_memory < 100 * 1024 * 1024:  # If less than 100 MB is free
+        print("Low GPU memory, switching to CPU.")
+        device = "cpu"
+    else:
+        device = "cuda"
+else:
+    device = "cpu"
 
-# 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"
+dtype = torch.float16 if device == "cuda" else torch.float32  # Use float16 for GPU and float32 for CPU
+
+# Clear any existing GPU memory cache
+torch.cuda.empty_cache()
 
-# Load the model with a smaller precision
+# Load the pipeline
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
-# Reduce inference steps to save memory
-def generate_image_flux(prompt, seed=42, width=width, height=height, num_inference_steps=2):
+# Reduce the inference steps and image dimensions
+def generate_image_flux(prompt, seed=42, width=512, height=512, num_inference_steps=2):
     generator = torch.Generator(device).manual_seed(seed)
     image = pipe(
         prompt=prompt,
         width=width,
         height=height,
-        num_inference_steps=num_inference_steps,  # Reduce the number of inference steps
+        num_inference_steps=num_inference_steps,  # Reduce steps to save memory
         generator=generator,
         guidance_scale=0.0
     ).images[0]