Update stablediffusion.py

stablediffusion.py

@@ -99,7 +99,7 @@ def generate_image(prompt,concept_embed,num_inference_steps=50,color_postprocess
         latents = latents * scheduler.init_noise_sigma # Scaling (previous versions did latents = latents * self.scheduler.sigmas[0]
 
       # Loop
-        with autocast("cpu"):  # will fallback to CPU if no CUDA; no autocast for MPS
+        with autocast("cuda"):  # will fallback to CPU if no CUDA; no autocast for MPS
             for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):
               # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
                     latent_model_input = torch.cat([latents] * 2)