Update app.py

app.py

@@ -47,10 +47,8 @@ if torch.cuda.is_available():
         previewer = Previewer()
         previewer_state_dict = torch.load("previewer/previewer_v1_100k.pt", map_location=torch.device('cpu'))["state_dict"]
         previewer.load_state_dict(previewer_state_dict)
-        def callback_prior(self_remote, i, t, kwargs):
-            latents = kwargs["latents"]
+        def callback_prior(i, t, latents):
             output = previewer(latents)
-            print(output)
             output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy())
             return output
         callback_steps = 1
@@ -100,7 +98,8 @@ def generate(
         guidance_scale=prior_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
         generator=generator,
-        callback_on_step_end=callback_prior,
+        callback=callback_prior,
+        callback_steps=callback_steps
     )
 
     if PREVIEW_IMAGES: