Add application file

app.py

@@ -4,12 +4,12 @@ import torch
 from PIL import ImageEnhance, Image
 import numpy as np
 
-# Load Stable Diffusion pipeline
+# Load Stable Diffusion pipeline with torch.float16 for GPU, fallback to float32 for CPU
 model_id = "CompVis/stable-diffusion-v1-4"
 default_scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
-pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=default_scheduler, torch_dtype=torch.float16)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = pipe.to(device)
+torch_dtype = torch.float16 if device == "cuda" else torch.float32
+pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=default_scheduler, torch_dtype=torch_dtype).to(device)
 
 # Scheduler options
 schedulers = {