Update app.py

app.py

@@ -3,11 +3,12 @@
 import gradio as gr
 import torch
 from diffusers import FluxPipeline
-import os # <-- Import the os library to access secrets
+import os
 
 # --- 1. Model Configuration ---
-# Define the device, base model, and your LoRA adapter from the Hub
-device = "cuda"
+# Automatically detect if a GPU is available, otherwise use CPU
+# The app will only work well on a GPU, but this prevents it from crashing on startup.
+device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model_id = "black-forest-labs/FLUX.1-Krea-dev"
 lora_model_id = "tusharmagar/flux1-krea-dev-lora-solarpunk"
 
@@ -17,40 +18,36 @@ lora_model_id = "tusharmagar/flux1-krea-dev-lora-solarpunk"
 pipe = FluxPipeline.from_pretrained(
     base_model_id,
     torch_dtype=torch.bfloat16,
-    token=os.environ.get("HF_TOKEN") # <-- Use the token from secrets
+    token=os.environ.get("HF_TOKEN")
 )
-# Move the pipeline to the GPU
+# Move the pipeline to the selected device (GPU)
 pipe.to(device)
 
 # Load and fuse your LoRA weights into the pipeline.
-# The model will now use your fine-tuned adjustments for all subsequent inferences.
 pipe.load_lora_weights(lora_model_id)
 pipe.fuse_lora()
 
 
 # --- 3. Define the Inference Function ---
-# This function takes a user's prompt and returns a generated image.
-# We've simplified it by hardcoding optimal settings for this model.
 def generate_image(prompt, progress=gr.Progress(track_tqdm=True)):
     """
     Generates an image based on the provided prompt.
     """
-    print(f"Generating image for prompt: {prompt}")  # For debugging in the Space logs
+    print(f"Generating image for prompt: {prompt}")
+    
+    # Ensure the generator is on the correct device
+    generator = torch.Generator(device=device).manual_seed(42)
     
-    # Run the pipeline to generate the image.
-    # We use a generator for reproducibility.
-    # FLUX models work well with a low number of steps and no guidance scale.
     image = pipe(
         prompt=prompt,
         num_inference_steps=8,
         guidance_scale=0.0,
-        generator=torch.Generator(device=device).manual_seed(42)
+        generator=generator
     ).images[0]
     
     return image
 
 # --- 4. Create the Gradio Interface ---
-# This section builds the web user interface.
 css = """
 #col-container {
     margin: 0 auto;
@@ -59,17 +56,15 @@ css = """
 """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        # Add a title and description for your demo
         gr.Markdown(
             """
             # FLUX.1 Solarpunk LoRA ☀️
             An interactive demo for the [flux1-krea-dev-lora-solarpunk](https://huggingface.co/tusharmagar/flux1-krea-dev-lora-solarpunk) model.
-            This LoRA excels at creating dreamy, solarpunk imaginations of real-world cities. 
+            This LoRA excels at creating dreamy, solarpunk imaginations of real-world cities.
             **Don't forget to add the trigger word `[SLRPNK]` to your prompt!**
             """
         )
 
-        # Create the main input and output components
         with gr.Row():
             prompt_input = gr.Textbox(
                 label="Enter your prompt",
@@ -83,7 +78,6 @@ with gr.Blocks(css=css) as demo:
         
         image_output = gr.Image(label="Result", show_label=False)
 
-        # Add your specific examples to make it easy for users to try
         gr.Examples(
             examples=[
                 "Solarpunk London with hexagonal solar panels, white architecture, keeping Big Ben unchanged, with a double decker bus on the road [SLRPNK]",
@@ -96,7 +90,6 @@ with gr.Blocks(css=css) as demo:
             fn=generate_image,
         )
     
-    # Connect the button click and prompt submission to the generation function
     generate_button.click(
         fn=generate_image,
         inputs=[prompt_input],