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,15 +711,20 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 #     image_3 = generate_image(hardcoded_prompt_3)
 #     return image_1, image_2, image_3
 
+import gradio as gr
+import torch
 from diffusers import FluxPipeline
 
+# 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
 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
+    pipe.to(device)  # Move the model to the correct device (GPU/CPU)
     return pipe
 
-# Function to generate image using Flux bot
+# Function to generate image using Flux bot on the specified device
 def generate_image_flux(prompt):
     pipe = initialize_flux_bot()
     image = pipe(
@@ -727,7 +732,7 @@ def generate_image_flux(prompt):
         guidance_scale=0.0,
         num_inference_steps=4,
         max_sequence_length=256,
-        generator=torch.Generator("cpu").manual_seed(0)
+        generator=torch.Generator(device).manual_seed(0)
     ).images[0]
     return image