Update app.py

app.py

@@ -2,19 +2,18 @@ import streamlit as st
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL
 import torch
 
-# Load the model and scheduler
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 model_id = "sam749/Photon-v1"
-vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16).to("cuda")
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
 pipe.vae = vae
 pipe.to("cuda")
 
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 
-# Streamlit interface
 st.title("Text-to-Image Generator")
 
-# Prompt inputs
 prompt = st.text_input("Enter the prompt for the image:")
 negative_prompt = "cartoon, painting, illustration, (worst quality, low quality, normal quality:2)"