Update app.py

app.py

@@ -6,8 +6,13 @@ app = Flask(__name__)
 
 # Load the model
 model_id = "ZB-Tech/Text-to-Image"
-pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Use float32 if running on CPU
+torch_dtype = torch.float16 if device == "cuda" else torch.float32
+
+pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+pipeline.to(device)
 
 @app.route("/generate", methods=["POST"])
 def generate_image():