Update app.py

app.py

@@ -1125,47 +1125,72 @@ 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
-import os
-
-# 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"
+# import gradio as gr
+# import torch
+# from diffusers import FluxPipeline
+# import os
+
+# # 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 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=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
 
-# 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=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
+# Function to call the Flux API and generate images
+def generate_image_flux(prompt):
+    client = Client("black-forest-labs/FLUX.1-schnell")
+    result = client.predict(
+        prompt=prompt,
+        seed=0,
+        randomize_seed=True,
+        width=400,
+        height=400,
+        num_inference_steps=4,
+        api_name="/infer"
+    )
+    return result
+
+# Function to update images with the three prompts
 def update_images():
     image_1 = generate_image_flux(hardcoded_prompt_1)
     image_2 = generate_image_flux(hardcoded_prompt_2)
@@ -1177,7 +1202,6 @@ def update_images():
 
 
 
-
 def format_restaurant_hotel_info(name, link, location, phone, rating, reviews, snippet):
     return f"""
     {name}
@@ -1473,14 +1497,22 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
             events_output = gr.HTML(value=fetch_local_events())
 
         with gr.Column():
-            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)
+            # 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])
+          
+            # Displaying the images generated using Flux API
+            image_output_1 = gr.Image(label="Image 1", elem_id="flux_image_1", width=400, height=400)
+            image_output_2 = gr.Image(label="Image 2", elem_id="flux_image_2", width=400, height=400)
+            image_output_3 = gr.Image(label="Image 3", elem_id="flux_image_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])
-