Update app.py

app.py

@@ -1,22 +1,38 @@
 import gradio as gr
 import torch
-from diffusers import FluxPipeline
+from diffusers import DiffusionPipeline  # Note: Change `FluxPipeline` to `DiffusionPipeline` if `FluxPipeline` is not correct
 from PIL import Image
 
-# Load the diffusion model
-pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-pipeline.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+# Function to determine the device and handle model loading
+def setup_pipeline():
+    # Check for CUDA availability
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Load the diffusion model
+    try:
+        pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+        if device == "cpu":
+            # If using CPU, ensure model is offloaded to avoid GPU-specific features
+            pipeline.enable_model_cpu_offload()
+        else:
+            # Move model to GPU
+            pipeline.to(device)
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        raise e
 
+    return pipeline, device
 
-# Set the model to the appropriate device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipeline.to(device)
+pipeline, device = setup_pipeline()
 
 def generate_image(prompt, guidance_scale=7.5, num_inference_steps=50):
     # Generate an image based on the prompt
     with torch.no_grad():
-        # Generate images
-        images = pipeline(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
+        try:
+            images = pipeline(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
+        except Exception as e:
+            print(f"Error generating image: {e}")
+            raise e
     
     # Assuming pipeline returns a list of images, just take the first one
     img = images[0]