Update app.py

app.py

@@ -8,6 +8,8 @@ import gradio as gr
 import torch
 from diffusers import StableDiffusionPipeline
 from torch import autocast
+
+from contextlib import nullcontext
 #from PIL import Image
 #from torchvision import transforms
 
@@ -18,6 +20,7 @@ openai.api_key = os.getenv('openaikey')
 authtoken = os.getenv('authtoken')
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+context = autocast if device == "cuda" else nullcontext
 pipe = StableDiffusionPipeline.from_pretrained("stale2000/sd-dnditem", torch_dtype=torch.float16, use_auth_token=authtoken)  
 pipe = pipe.to(device)
 
@@ -55,9 +58,13 @@ def predict(input, manual_query_repacement, history=[]):
     #      return images, False
     #  pipe.safety_checker = null_safety
 
-    with autocast("cuda"):
-      images = pipe(n_samples*[prompt], guidance_scale=scale).images
+    
+    #with autocast("cuda"):
+    #  images = pipe(n_samples*[prompt], guidance_scale=scale).images
 
+    with context("cuda"):
+        images = pipe(n_samples*[prompt], guidance_scale=scale, num_inference_steps=5).images
+    
     for idx, im in enumerate(images):
       im.save(f"{idx:06}.png")