Update app.py

app.py

@@ -713,29 +713,53 @@ def generate_map(location_names):
 #     return image_1, image_2, image_3
 
 
+import gradio as gr
 import torch
 from diffusers import FluxPipeline
+import os
 
-def initialize_flux_pipeline():
-    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-    pipe.enable_model_cpu_offload()  # Offload to CPU to save VRAM
-    return pipe
+# Set PYTORCH_CUDA_ALLOC_CONF to handle memory fragmentation
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 
-# Initialize the model
-flux_pipe = initialize_flux_pipeline()
+# Check if CUDA (GPU) is available, otherwise fallback to CPU
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
+# 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
 
-def generate_flux_image(prompt):
-    # Use the initialized flux_pipe to generate an image based on the input prompt
-    image = flux_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=torch.Generator("cpu").manual_seed(0)
+        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]
     return image
 
+# 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"
+
+# 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
+
 
 
 
@@ -1481,20 +1505,20 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
 
         with gr.Column():
 
-            flux_prompt = gr.Textbox(show_copy_button=True, label="Flux Prompt", placeholder="Enter prompt for Flux image generation")
-            flux_image_output = gr.Image()
-            flux_generate_button = gr.Button("Generate Flux Image")
-            
-            # When the button is clicked, the image generation function is triggered
-            flux_generate_button.click(fn=generate_flux_image, inputs=flux_prompt, outputs=flux_image_output)
+            # Display images
+            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])
 
 
 
 
-            # 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])
+
+           
 
 demo.queue()
 demo.launch(show_error=True)