Update app.py

app.py

@@ -6,10 +6,10 @@ import torch
 def load_model(base_model_id, adapter_model_id=None):
     if torch.cuda.is_available():
         device = "cuda"
-        info = "Running on GPU (CUDA)"
+        info = "Running on GPU (CUDA) 🔥"
     else:
         device = "cpu"
-        info = "Running on CPU"
+        info = "Running on CPU 🥶"
 
     # Load the base model dynamically on the correct device
     pipe = StableDiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
@@ -23,24 +23,25 @@ def load_model(base_model_id, adapter_model_id=None):
     
     return pipe, info
 
-
-if torch.cuda.is_available():
-        device = "cuda"
-        info = "Running on GPU (CUDA) 🔥"
-    else:
-        device = "cpu"
-        info = "Running on CPU 🥶"
-
 # Function for text-to-image generation with dynamic model ID and device info
 def generate_image(base_model_id, adapter_model_id, prompt):
     pipe, info = load_model(base_model_id, adapter_model_id)
     image = pipe(prompt).images[0]
     return image, info
 
+# Check device (GPU/CPU) once at the start and show it in the UI
+if torch.cuda.is_available():
+    device = "cuda"
+    info = "Running on GPU (CUDA) 🔥"
+else:
+    device = "cpu"
+    info = "Running on CPU 🥶"
+
 # Create the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("## Custom Text-to-Image Generator with Adapter Support")
-    gr.Markdown(f"{info}")
+    gr.Markdown(f"**{info}**")  # Display GPU/CPU information in the UI
+
     with gr.Row():
         with gr.Column():
             base_model_id = gr.Textbox(label="Enter Base Model ID (e.g., CompVis/stable-diffusion-v1-4)", placeholder="Base Model ID")