Update app.py

app.py

@@ -692,41 +692,6 @@ def generate_map(location_names):
 from diffusers import DiffusionPipeline
 import torch
 
-# Load the FLUX model
-def load_flux_model():
-    pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell")
-    if torch.cuda.is_available():
-        pipeline = pipeline.to("cuda")
-    return pipeline
-
-flux_pipeline = load_flux_model()
-
-# Function to generate an image based on a selected prompt
-def generate_flux_image(selected_prompt):
-    prompt_map = {
-        "Toyota Truck in Birmingham": "A high quality cinematic image for Toyota Truck in Birmingham skyline shot in the style of Michael Mann",
-        "Alabama Quarterback": "A high quality cinematic image for Alabama Quarterback close up emotional shot in the style of Michael Mann",
-        "Taylor Swift Concert": "A high quality cinematic image for Taylor Swift concert in Birmingham skyline style of Michael Mann"
-    }
-    prompt = prompt_map.get(selected_prompt, "")
-    if not prompt:
-        return None  # In case no prompt is matched
-
-    # Generate the image using the FLUX model
-    image = flux_pipeline(prompt).images[0]
-    
-    # Save the image to a temporary file
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
-        image.save(f.name)
-        return f.name
-
-# Predefine the three prompts for automatic image generation
-default_prompts = [
-    "Toyota Truck in Birmingham",
-    "Alabama Quarterback",
-    "Taylor Swift Concert"
-]
-
 
 
 
@@ -1160,7 +1125,44 @@ def handle_model_choice_change(selected_model):
         # Default case: allow interaction
         return gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True)
 
+import gradio as gr
+import torch
+from diffusers import FluxPipeline
+from PIL import Image
+
+# Flux Image Generation
+# 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
+
+# Function to generate image using Flux
+def generate_flux_image(prompt: str):
+    generator = torch.Generator("cpu").manual_seed(0)  # For reproducibility
+    image = flux_pipe(
+        prompt,
+        guidance_scale=0.0,
+        num_inference_steps=4,
+        max_sequence_length=256,
+        generator=generator
+    ).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
+
+# 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"
 
+# Generate the images automatically upon launching the interface
+def show_images_on_startup():
+    img1 = generate_flux_image(hardcoded_prompt_1)
+    img2 = generate_flux_image(hardcoded_prompt_2)
+    img3 = generate_flux_image(hardcoded_prompt_3)
+    return img1, img2, img3
 
 
 
@@ -1461,23 +1463,11 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
             news_output = gr.HTML(value=fetch_local_news())
             events_output = gr.HTML(value=fetch_local_events())
 
-
         with gr.Column():
-
-            # Image output section
-            gr.Markdown("<h2>Generated Images on Load</h2>")
-            image_output_1 = gr.Image(type="filepath", label="Generated Image 1", width=400, height=400)
-            image_output_2 = gr.Image(type="filepath", label="Generated Image 2", width=400, height=400)
-            image_output_3 = gr.Image(type="filepath", label="Generated Image 3", width=400, height=400)
-
-            # Automatically generate and display the three images on startup using the predefined prompts
-            startup_image_1 = generate_flux_image(default_prompts[0])
-            startup_image_2 = generate_flux_image(default_prompts[1])
-            startup_image_3 = generate_flux_image(default_prompts[2])
-
-            image_output_1.update(value=startup_image_1)
-            image_output_2.update(value=startup_image_2)
-            image_output_3.update(value=startup_image_3)
+            gr.Markdown("<h1 style='color: blue;'>Cinematic Images Generated with Flux</h1>")
+            image_output_1 = gr.Image()
+            image_output_2 = gr.Image()
+            image_output_3 = gr.Image()