Update app.py

app.py

@@ -24,7 +24,7 @@ CACHE_EXAMPLES = False #torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1536"))
 USE_TORCH_COMPILE = False
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
-PREVIEW_IMAGES = False
+PREVIEW_IMAGES = True
 
 dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -47,10 +47,12 @@ 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(i, t, latents):
+        def callback_prior(pipeline, step_index, t, callback_kwargs):
+            latents = callback_kwargs["latents"]
             output = previewer(latents)
             output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy())
-            return output
+            callback_kwargs["preview_output"] = output
+            return callback_kwargs
         callback_steps = 1
     else:
         previewer = None
@@ -84,7 +86,7 @@ def generate(
     profile: gr.OAuthProfile | None = None,
 ) -> PIL.Image.Image:
     
-    #previewer.eval().requires_grad_(False).to(device).to(dtype)
+    previewer.eval().requires_grad_(False).to(device).to(dtype)
     prior_pipeline.to(device)
     decoder_pipeline.to(device)
     
@@ -100,8 +102,8 @@ def generate(
         guidance_scale=prior_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
         generator=generator,
-        #callback=callback_prior,
-        #callback_steps=callback_steps
+        callback_on_step_end=callback_prior,
+        callback_on_step_end_tensor_inputs=['latents']
     )
     print(prior_output)
     if PREVIEW_IMAGES: