debugging silent crash

app.py

@@ -68,6 +68,8 @@ def get_inference_lambda(seed):
 
     image_width = image_height = 256
 
+    print("all models setup")
+
     def __tokenize_prompt(prompt: str):
 
         return tokenizer(
@@ -79,10 +81,12 @@ def get_inference_lambda(seed):
         ).input_ids.astype(jnp.float32)
 
     def __convert_image(vae_output):
-        return [
-            Image.fromarray(image)
-            for image in (np.asarray(vae_output) * 255).round().astype(np.uint8)
-        ]
+        print("skipping image conversion...")
+        return None
+        # return [
+        #     Image.fromarray(image)
+        #     for image in (np.asarray(vae_output) * 255).round().astype(np.uint8)
+        # ]
 
     def __predict_image(tokenized_prompt: jnp.array):
 
@@ -92,10 +96,10 @@ def get_inference_lambda(seed):
             params=text_encoder_params,
             train=False,
         )[0]
-
         context = jnp.concatenate(
             [negative_prompt_text_encoder_hidden_states, text_encoder_hidden_states]
         )
+        jax.debug.print("got text encoding...")
 
         latent_shape = (
             tokenized_prompt.shape[0],
@@ -152,6 +156,7 @@ def get_inference_lambda(seed):
         initial_scheduler_state = scheduler.set_timesteps(
             scheduler.create_state(), num_inference_steps=timesteps, shape=latent_shape
         )
+        jax.debug.print("initialized scheduler state...")
 
         # initialize latents
         initial_latents = (
@@ -160,11 +165,11 @@ def get_inference_lambda(seed):
             )
             * initial_scheduler_state.init_noise_sigma
         )
+        jax.debug.print("initialized latents...")
 
         final_latents, _ = jax.lax.fori_loop(
             0, timesteps, ___timestep, (initial_latents, initial_scheduler_state)
         )
-
         jax.debug.print("got final latents...")
 
         # scale and decode the image latents with vae
@@ -181,8 +186,7 @@ def get_inference_lambda(seed):
             .clip(0, 1)
             .transpose(0, 2, 3, 1)
         )
-
-        jax.debug.print("got vae decoded image output...")
+        jax.debug.print("got vae processed image output...")
 
         # return reshaped vae outputs
         return image
@@ -212,7 +216,7 @@ with gr.Blocks(theme="gradio/soft") as demo:
     with gr.Tab("Journal"):
         gr.Markdown(
             """
-            ## On How Four Crazy Fellows Embarked on Training a U-Net from Scratch in Five Days with JAX and Almost Died in the End
+            ## On How Four Crazy Fellows Embarked on Training a JAX U-Net from Scratch in Five Days and Almost Died in the End
 
             Lorem ipsum dolor sit amet, consectetur adipiscing elit. Mauris vitae varius libero. Nullam laoreet eget sapien quis tristique. Cras odio odio, consequat sed cursus quis, dignissim hendrerit ligula. Curabitur non lorem tellus. Nam bibendum malesuada mi sed faucibus. Sed euismod enim metus, sit amet venenatis elit elementum vel. Duis nec rhoncus tellus, rhoncus auctor justo. Proin id gravida dolor. Sed nulla lectus, finibus non fringilla ac, fermentum in sapien. Cras lobortis est augue, vel posuere justo pretium vitae. Aliquam lorem dolor, condimentum et finibus rutrum, rhoncus eget nunc.