Update app.py

app.py

@@ -92,77 +92,70 @@ def edit(input_image,
     # invert and retrieve noise maps and latent
     wt, zs, wts = invert(x0 =x0 , prompt_src=src_prompt, num_diffusion_steps=steps, cfg_scale_src=cfg_scale_src)
 
-    #
-    xT=wts[skip]
-    etas=1.0
-    prompts=[tar_prompt]
-    cfg_scales=[cfg_scale_tar]
-    prog_bar=False
-    zs=zs[skip:]
+    # #
+    # xT=wts[skip]
+    # etas=1.0
+    # prompts=[tar_prompt]
+    # cfg_scales=[cfg_scale_tar]
+    # prog_bar=False
+    # zs=zs[skip:]
 
 
-    batch_size = len(prompts)
+    # batch_size = len(prompts)
 
-    cfg_scales_tensor = torch.Tensor(cfg_scales).view(-1,1,1,1).to(sd_pipe.device)
+    # cfg_scales_tensor = torch.Tensor(cfg_scales).view(-1,1,1,1).to(sd_pipe.device)
 
-    text_embeddings = encode_text(sd_pipe, prompts)
-    uncond_embedding = encode_text(sd_pipe, [""] * batch_size)
+    # text_embeddings = encode_text(sd_pipe, prompts)
+    # uncond_embedding = encode_text(sd_pipe, [""] * batch_size)
 
-    if etas is None: etas = 0
-    if type(etas) in [int, float]: etas = [etas]*sd_pipe.scheduler.num_inference_steps
-    assert len(etas) == sd_pipe.scheduler.num_inference_steps
-    timesteps = sd_pipe.scheduler.timesteps.to(sd_pipe.device)
+    # if etas is None: etas = 0
+    # if type(etas) in [int, float]: etas = [etas]*sd_pipe.scheduler.num_inference_steps
+    # assert len(etas) == sd_pipe.scheduler.num_inference_steps
+    # timesteps = sd_pipe.scheduler.timesteps.to(sd_pipe.device)
 
-    xt = xT.expand(batch_size, -1, -1, -1)
-    op = tqdm(timesteps[-zs.shape[0]:]) if prog_bar else timesteps[-zs.shape[0]:] 
+    # xt = xT.expand(batch_size, -1, -1, -1)
+    # op = tqdm(timesteps[-zs.shape[0]:]) if prog_bar else timesteps[-zs.shape[0]:] 
 
-    t_to_idx = {int(v):k for k,v in enumerate(timesteps[-zs.shape[0]:])}
+    # t_to_idx = {int(v):k for k,v in enumerate(timesteps[-zs.shape[0]:])}
 
-    for t in op:
-        idx = t_to_idx[int(t)]        
-        ## Unconditional embedding
-        with torch.no_grad():
-            uncond_out = sd_pipe.unet.forward(xt, timestep =  t, 
-                                            encoder_hidden_states = uncond_embedding)
+    # for t in op:
+    #     idx = t_to_idx[int(t)]        
+    #     ## Unconditional embedding
+    #     with torch.no_grad():
+    #         uncond_out = sd_pipe.unet.forward(xt, timestep =  t, 
+    #                                         encoder_hidden_states = uncond_embedding)
 
-            ## Conditional embedding  
-        if prompts:  
-            with torch.no_grad():
-                cond_out = sd_pipe.unet.forward(xt, timestep =  t, 
-                                                encoder_hidden_states = text_embeddings)
+    #         ## Conditional embedding  
+    #     if prompts:  
+    #         with torch.no_grad():
+    #             cond_out = sd_pipe.unet.forward(xt, timestep =  t, 
+    #                                             encoder_hidden_states = text_embeddings)
             
         
-        z = zs[idx] if not zs is None else None
-        z = z.expand(batch_size, -1, -1, -1)
-        if prompts:
-            ## classifier free guidance
-            noise_pred = uncond_out.sample + cfg_scales_tensor * (cond_out.sample - uncond_out.sample)
-        else: 
-            noise_pred = uncond_out.sample
-        # 2. compute less noisy image and set x_t -> x_t-1  
-        xt = reverse_step(sd_pipe, noise_pred, t, xt, eta = etas[idx], variance_noise = z)
-
-        # interm denoised img
-        with autocast("cuda"), inference_mode():
-            x0_dec = sd_pipe.vae.decode(1 / 0.18215 * xt).sample
-            if x0_dec.dim()<4:
-                x0_dec = x0_dec[None,:,:,:]
-            interm_img = image_grid(x0_dec)
-            yield interm_img
+    #     z = zs[idx] if not zs is None else None
+    #     z = z.expand(batch_size, -1, -1, -1)
+    #     if prompts:
+    #         ## classifier free guidance
+    #         noise_pred = uncond_out.sample + cfg_scales_tensor * (cond_out.sample - uncond_out.sample)
+    #     else: 
+    #         noise_pred = uncond_out.sample
+    #     # 2. compute less noisy image and set x_t -> x_t-1  
+    #     xt = reverse_step(sd_pipe, noise_pred, t, xt, eta = etas[idx], variance_noise = z)
+
+    #     # interm denoised img
+    #     with autocast("cuda"), inference_mode():
+    #         x0_dec = sd_pipe.vae.decode(1 / 0.18215 * xt).sample
+    #         if x0_dec.dim()<4:
+    #             x0_dec = x0_dec[None,:,:,:]
+    #         interm_img = image_grid(x0_dec)
+    #         yield interm_img
       
-    yield interm_img
+    # yield interm_img
     
-    # # vae decode image
-    # with autocast("cuda"), inference_mode():
-    #     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)
-    # return img
 
-    # output = sample(wt, zs, wts, prompt_tar=tar_prompt)
+    output = sample(wt, zs, wts, prompt_tar=tar_prompt, cfg_scale_tar=cfg_scale_tar, skip=skip)
 
-    # return output
+    return output