Update app.py

app.py

@@ -688,11 +688,49 @@ def generate_map(location_names):
     map_html = m._repr_html_()
     return map_html
 
-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)
-
-def generate_image(prompt):
+# 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)
+
+# def generate_image(prompt):
+#     with torch.cuda.amp.autocast():
+#         image = pipe(
+#             prompt,
+#             num_inference_steps=28,
+#             guidance_scale=3.0,
+#         ).images[0]
+#     return image
+
+# hardcoded_prompt_1 = "A high quality cinematic image for Toyota Truck in Birmingham skyline shot in th style of Michael Mann"
+# hardcoded_prompt_2 = "A high quality cinematic image for Alabama Quarterback close up emotional shot in th style of Michael Mann"
+# hardcoded_prompt_3 = "A high quality cinematic image for Taylor Swift concert in Birmingham skyline style of Michael Mann"
+
+# def update_images():
+#     image_1 = generate_image(hardcoded_prompt_1)
+#     image_2 = generate_image(hardcoded_prompt_2)
+#     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
+
+# Function to initialize Flux bot model
+def initialize_flux_bot():
+    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+    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,
@@ -701,14 +739,35 @@ def generate_image(prompt):
         ).images[0]
     return image
 
-hardcoded_prompt_1 = "A high quality cinematic image for Toyota Truck in Birmingham skyline shot in th style of Michael Mann"
-hardcoded_prompt_2 = "A high quality cinematic image for Alabama Quarterback close up emotional shot in th style of Michael Mann"
+# Function to generate image using Flux bot
+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)
+    ).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"
 
-def update_images():
-    image_1 = generate_image(hardcoded_prompt_1)
-    image_2 = generate_image(hardcoded_prompt_2)
-    image_3 = generate_image(hardcoded_prompt_3)
+# 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)
     return image_1, image_2, image_3
 
 
@@ -1369,6 +1428,8 @@ 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")
@@ -1449,9 +1510,14 @@ 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(hardcoded_prompt_1), width=400, height=400)
-            image_output_2 = gr.Image(value=generate_image(hardcoded_prompt_2), width=400, height=400)
-            image_output_3 = gr.Image(value=generate_image(hardcoded_prompt_3), width=400, height=400)
+            # image_output_1 = gr.Image(value=generate_image(hardcoded_prompt_1), width=400, height=400)
+            # image_output_2 = gr.Image(value=generate_image(hardcoded_prompt_2), width=400, height=400)
+            # 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)
 
             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")