Update app.py

app.py

@@ -1128,38 +1128,49 @@ def handle_model_choice_change(selected_model):
 import gradio as gr
 import torch
 from diffusers import FluxPipeline
-from PIL import Image
+import os
 
-# Load the Flux pipeline
-flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-flux_pipe.enable_model_cpu_offload()  # Save some VRAM by offloading to CPU if needed
+# Set PYTORCH_CUDA_ALLOC_CONF to handle memory fragmentation
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+
+# Check if CUDA (GPU) is available, otherwise fallback to CPU
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# Function to generate image using Flux
-def generate_flux_image(prompt: str):
-    generator = torch.Generator("cpu").manual_seed(0)  # For reproducibility
-    image = flux_pipe(
+# Function to initialize Flux bot model with GPU memory management
+def initialize_flux_bot():
+    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
+def generate_image_flux(prompt):
+    pipe = initialize_flux_bot()
+    image = pipe(
         prompt,
         guidance_scale=0.0,
-        num_inference_steps=4,
-        max_sequence_length=256,
-        generator=generator
+        num_inference_steps=2,  # Reduced steps to save memory
+        max_sequence_length=128,  # Reduced sequence length to save memory
+        generator=torch.Generator(device).manual_seed(0)
     ).images[0]
-    
-    # Save image temporarily and return for display
-    temp_image_path = f"temp_flux_image_{hash(prompt)}.png"
-    image.save(temp_image_path)
-    
-    return temp_image_path
+    return image
 
-# Hardcoded prompts for generating images
+# Hardcoded prompts for the images
 hardcoded_prompt_1 = "A high quality cinematic image for Toyota Truck in Birmingham skyline shot in the style of Michael Mann"
 hardcoded_prompt_2 = "A high quality cinematic image for Alabama Quarterback close up emotional shot in the style of Michael Mann"
 hardcoded_prompt_3 = "A high quality cinematic image for Taylor Swift concert in Birmingham skyline style of Michael Mann"
 
-# Generate the images immediately
-img1_path = generate_flux_image(hardcoded_prompt_1)
-# img2_path = generate_flux_image(hardcoded_prompt_2)
-# img3_path = generate_flux_image(hardcoded_prompt_3)
+# Function to update images
+def update_images():
+    image_1 = generate_image_flux(hardcoded_prompt_1)
+    image_2 = generate_image_flux(hardcoded_prompt_2)
+    image_3 = generate_image_flux(hardcoded_prompt_3)
+    return image_1, image_2, image_3
 
 
 
@@ -1462,10 +1473,13 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
             events_output = gr.HTML(value=fetch_local_events())
 
         with gr.Column():
-            # Display the pre-generated images directly
-            image_output_1 = gr.Image(value=img1_path)
-            # image_output_2 = gr.Image(value=img2_path)
-            # image_output_3 = gr.Image(value=img3_path)
+            image_output_1 = gr.Image(value=generate_image_flux(hardcoded_prompt_1), width=400, height=400)
+            image_output_2 = gr.Image(value=generate_image_flux(hardcoded_prompt_2), width=400, height=400)
+            image_output_3 = gr.Image(value=generate_image_flux(hardcoded_prompt_3), width=400, height=400)
+    
+            # Refresh button to update images
+            refresh_button = gr.Button("Refresh Images")
+            refresh_button.click(fn=update_images, inputs=None, outputs=[image_output_1, image_output_2, image_output_3])