Update app.py

app.py

@@ -20,11 +20,11 @@ if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶</p>"
 
 MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") != "0"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
+CACHE_EXAMPLES = False #torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") != "0"
+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 = True
+PREVIEW_IMAGES = False
 
 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
@@ -62,6 +64,7 @@ else:
 
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    print("randomizing seed")
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
@@ -82,7 +85,8 @@ def generate(
     num_images_per_prompt: int = 2,
     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)
     
@@ -98,10 +102,9 @@ 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']
     )
-
     if PREVIEW_IMAGES:
         for _ in range(len(DEFAULT_STAGE_C_TIMESTEPS)):
             r = next(prior_output)
@@ -119,7 +122,7 @@ def generate(
         generator=generator,
         output_type="pil",
     ).images
-
+    print(decoder_output)
     #Save images
     for image in decoder_output:
         user_history.save_image(
@@ -137,14 +140,14 @@ def generate(
                 "num_images_per_prompt": num_images_per_prompt,
             },
         )
-
+    
     yield decoder_output[0]
 
 
 examples = [
     "An astronaut riding a green horse",
     "A mecha robot in a favela by Tarsila do Amaral",
-    "The sprirt of a Tamagotchi wandering in the city of Los Angeles",
+    "The spirit of a Tamagotchi wandering in the city of Los Angeles",
     "A delicious feijoada ramen dish"
 ]
 
@@ -186,12 +189,14 @@ with gr.Blocks() as demo:
                 label="Width",
                 minimum=1024,
                 maximum=MAX_IMAGE_SIZE,
+                step=512,
                 value=1024,
             )
             height = gr.Slider(
                 label="Height",
                 minimum=1024,
                 maximum=MAX_IMAGE_SIZE,
+                step=512,
                 value=1024,
             )
             num_images_per_prompt = gr.Slider(