Update app.py

app.py

@@ -688,7 +688,7 @@ def generate_map(location_names):
     map_html = m._repr_html_()
     return map_html
 
-# device = "cuda:0" if torch.cuda.is_available() else "cpu"
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 # pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16)
 # pipe.to(device)
 
@@ -711,16 +711,7 @@ def generate_map(location_names):
 #     image_3 = generate_image(hardcoded_prompt_3)
 #     return image_1, image_2, image_3
 
-
-
-from diffusers import StableDiffusionPipeline, FluxPipeline
-
-# Function to initialize Stable Diffusion model
-def initialize_stable_diffusion():
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16)
-    pipe.to(device)
-    return pipe
+from diffusers import FluxPipeline
 
 # Function to initialize Flux bot model
 def initialize_flux_bot():
@@ -728,17 +719,6 @@ def initialize_flux_bot():
     pipe.enable_model_cpu_offload()  # Saves VRAM by offloading the model to CPU
     return pipe
 
-# Function to generate image using Stable Diffusion
-def generate_image_stable_diffusion(prompt):
-    pipe = initialize_stable_diffusion()
-    with torch.cuda.amp.autocast():
-        image = pipe(
-            prompt,
-            num_inference_steps=28,
-            guidance_scale=3.0,
-        ).images[0]
-    return image
-
 # Function to generate image using Flux bot
 def generate_image_flux(prompt):
     pipe = initialize_flux_bot()
@@ -751,26 +731,22 @@ def generate_image_flux(prompt):
     ).images[0]
     return image
 
-# Combined function to handle model switching based on radio button selection
-def generate_image(prompt, model_choice):
-    if model_choice == "IG-1":
-        return generate_image_flux(prompt)
-    else:  # Default to Stable Diffusion
-        return generate_image_stable_diffusion(prompt)
-
 # 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 based on the selected model
-def update_images(model_choice):
-    image_1 = generate_image(hardcoded_prompt_1, model_choice)
-    image_2 = generate_image(hardcoded_prompt_2, model_choice)
-    image_3 = generate_image(hardcoded_prompt_3, model_choice)
+# 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
 
 
+
+
+
 def fetch_local_news():
     api_key = os.environ['SERP_API']
     url = f'https://serpapi.com/search.json?engine=google_news&q=birmingham headline&api_key={api_key}'
@@ -1428,8 +1404,7 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
 
             chat_input = gr.Textbox(show_copy_button=True, interactive=True, show_label=False, label="ASK Radar !!!", placeholder="Hey Radar...!!")
             tts_choice = gr.Radio(label="Select TTS System", choices=["Alpha", "Beta"], value="Alpha")
-            # Add a radio button to choose between IG-1 (Flux) and IG-2 (Stable Diffusion), defaulting to IG-1
-            model_choice = gr.Radio(label="Select Image Generation Model", choices=["IG-1", "IG-2"], value="IG-1")
+        
             retriever_button = gr.Button("Retriever")
 
             clear_button = gr.Button("Clear")
@@ -1515,12 +1490,13 @@ with gr.Blocks(theme='Pijush2023/scikit-learn-pijush') as demo:
             # image_output_3 = gr.Image(value=generate_image(hardcoded_prompt_3), width=400, height=400)
 
             # Display images
-            image_output_1 = gr.Image(value=generate_image(hardcoded_prompt_1, "IG-1"), width=400, height=400)
-            image_output_2 = gr.Image(value=generate_image(hardcoded_prompt_2, "IG-1"), width=400, height=400)
-            image_output_3 = gr.Image(value=generate_image(hardcoded_prompt_3, "IG-1"), 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], api_name="update_image")
+            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)