Update app.py

app.py

@@ -713,51 +713,52 @@ def generate_map(location_names):
 #     return image_1, image_2, image_3
 
 
-import gradio as gr
 import torch
-from diffusers import FluxPipeline
-import os
+from diffusers import DiffusionPipeline
+import numpy as np
+import random
+import gradio as gr
 
-# Set PYTORCH_CUDA_ALLOC_CONF to handle memory fragmentation
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+# Constants for device and dtype
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Check if CUDA (GPU) is available, otherwise fallback to CPU
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# Load the Flux pipeline model
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
-# 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()
+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,
-        guidance_scale=0.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)
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        guidance_scale=0.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"
+    
+    return image, seed
 
-# Function to update images
+# Existing image generation function, updated to use the new Flux pipeline
 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)
+    # Hardcoded prompts for generating 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"
+    
+    # Use the Flux pipeline to generate images based on hardcoded prompts
+    image_1, _ = infer(hardcoded_prompt_1)
+    image_2, _ = infer(hardcoded_prompt_2)
+    image_3, _ = infer(hardcoded_prompt_3)
+    
     return image_1, image_2, image_3
 
 
@@ -1505,7 +1506,7 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
 
         with gr.Column():
 
-            # Display images
+            # Display generated 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)