Update app.py

app.py

@@ -89,13 +89,14 @@ def predict(prompt, weights, seed, scale=7.5, steps=50):
     try:
         with torch.no_grad():
             has_cuda = torch.cuda.is_available()
-            if has_cuda:
-                generator = torch.Generator('cuda').manual_seed(int(seed))
-            else:
-                generator = torch.Generator().manual_seed(int(seed))
-            image_perpneg = pipe(prompt, guidance_scale=float(scale), generator=generator,
+            with autocast('cpu' if not has_cuda else 'cuda'):
+                if has_cuda:
+                    generator = torch.Generator('cuda').manual_seed(int(seed))
+                else:
+                    generator = torch.Generator().manual_seed(int(seed))
+                image_perpneg = pipe(prompt, guidance_scale=float(scale), generator=generator,
                             num_inference_steps=steps, weights=weights)["images"][0]
-            return image_perpneg
+                return image_perpneg
     except Exception as e:
         print(e)
         return None