Update app.py

app.py

@@ -22,8 +22,7 @@ def invert(x0, prompt_src="", num_diffusion_steps=100, cfg_scale_src = 3.5, eta
   sd_pipe.scheduler.set_timesteps(num_diffusion_steps)
 
   # vae encode image
-  with autocast("cuda"), inference_mode():
-      w0 = (sd_pipe.vae.encode(x0).latent_dist.mode() * 0.18215).float()
+  w0 = (sd_pipe.vae.encode(x0).latent_dist.mode() * 0.18215).float()
 
   # find Zs and wts - forward process
   wt, zs, wts = inversion_forward_process(sd_pipe, w0, etas=eta, prompt=prompt_src, cfg_scale=cfg_scale_src, prog_bar=True, num_inference_steps=num_diffusion_steps)
@@ -37,8 +36,7 @@ def sample(wt, zs, wts, prompt_tar="", cfg_scale_tar=15, skip=36, eta = 1):
     w0, _ = inversion_reverse_process(sd_pipe, xT=wts[skip], etas=eta, prompts=[prompt_tar], cfg_scales=[cfg_scale_tar], prog_bar=True, zs=zs[skip:])
     
     # vae decode image
-    with autocast("cuda"), inference_mode():
-        x0_dec = sd_pipe.vae.decode(1 / 0.18215 * w0).sample
+    x0_dec = sd_pipe.vae.decode(1 / 0.18215 * w0).sample
     if x0_dec.dim()<4:
         x0_dec = x0_dec[None,:,:,:]
     img = image_grid(x0_dec)
@@ -47,9 +45,9 @@ def sample(wt, zs, wts, prompt_tar="", cfg_scale_tar=15, skip=36, eta = 1):
 # load pipelines
 sd_model_id = "runwayml/stable-diffusion-v1-5"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-sd_pipe = StableDiffusionPipeline.from_pretrained(sd_model_id).to(device)
+sd_pipe = StableDiffusionPipeline.from_pretrained(sd_model_id, ,torch_dtype=torch.float16).to(device)
 sd_pipe.scheduler = DDIMScheduler.from_config(sd_model_id, subfolder = "scheduler")
-sem_pipe = SemanticStableDiffusionPipeline.from_pretrained(sd_model_id).to(device)
+sem_pipe = SemanticStableDiffusionPipeline.from_pretrained(sd_model_id, ,torch_dtype=torch.float16).to(device)
 
 
 def edit(input_image, input_image_prompt='', target_prompt='', edit_prompt='',