Update app.py

app.py

@@ -20,16 +20,24 @@ def generate_and_display_images(model_selection, scenery, style, height, width,
             return "Invalid seed value. Seed must be an integer."
     torch.manual_seed(seed)
 
+    # Check if CUDA is available and set the appropriate dtype
+    if torch.cuda.is_available():
+        device = "cuda"
+        dtype = torch.float16
+    else:
+        device = "cpu"
+        dtype = torch.float32
+
     prompt = f"Scenery: {scenery}; Style: {style}"
 
     generated_images = []
     if model_selection == "dreamlike-art/dreamlike-photoreal-2.0":
-        model = StableDiffusionPipeline.from_pretrained(model_selection, torch_dtype=torch.float16).to("cpu")
+        model = StableDiffusionPipeline.from_pretrained(model_selection, torch_dtype=dtype).to(device)
         for _ in range(num_images):
             image = model(prompt=prompt, num_inference_steps=n_steps, guidance_scale=guidance_scale, negative_prompt=negative_prompt, height=height, width=width).images[0]
             generated_images.append(image)
     else:
-        base = DiffusionPipeline.from_pretrained(model_selection, torch_dtype=torch.float16, use_auth_token=True).to("cpu")
+        base = DiffusionPipeline.from_pretrained(model_selection, torch_dtype=dtype, use_auth_token=True).to(device)
         for _ in range(num_images):
             if "refiner" in model_selection:
                 image_latent = base(prompt=prompt, num_inference_steps=n_steps, denoising_end=high_noise_frac, output_type="latent").images
@@ -37,7 +45,7 @@ def generate_and_display_images(model_selection, scenery, style, height, width,
             else:
                 image = base(prompt=prompt, num_inference_steps=n_steps, guidance_scale=guidance_scale, negative_prompt=negative_prompt, height=height, width=width).images[0]
             generated_images.append(image)
-    
+
     # Save images and return file paths for Gradio display
     file_paths = []
     for i, image in enumerate(generated_images):
@@ -48,8 +56,6 @@ def generate_and_display_images(model_selection, scenery, style, height, width,
 
     return file_paths
 
-
-
 # Define Gradio interface
 iface = gr.Interface(
     fn=generate_and_display_images,