Update app.py

app.py

@@ -17,6 +17,8 @@ def load_model(config_path, ckpt_path):
     with open(ckpt_path, "rb") as f:
         leaves = pickle.load(f)
 
+    leaves = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), leaves)
+
     from model import DiT, DiTConfig
 
     dit_config = DiTConfig(**config["model"])
@@ -32,13 +34,16 @@ def sample_images(graphdef, state, x0, t):
     flow = nnx.merge(graphdef, state)
 
     def flow_fn(y, t):
+        y = y.astype(jnp.bfloat16)
+        t = t.astype(jnp.bfloat16)
         o = flow(y, t[None])
-        return o
+        return o.astype(jnp.float32)
 
-    o = ode.odeint(flow_fn, x0, t, rtol=1e-2)
+    o = ode.odeint(flow_fn, x0, t, rtol=1e-4)
     o = jnp.clip(o[-1], 0, 1)
     return o
 
+
 @spaces.GPU
 def generate_grid(seed, noise_level):
     # Load model (doing this inside function to avoid global variables)
@@ -66,8 +71,6 @@ def generate_grid(seed, noise_level):
     return jax.device_get(grid)
 
 
-generate_grid(0, 3)
-
 # Create Gradio interface
 demo = gr.Interface(
     fn=generate_grid,