Update

app.py

@@ -92,7 +92,7 @@ def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmb
 def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
     return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, dtype=jnp.bfloat16)
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=60)
 def load_encoders():
     is_schnell = True
     t5 = load_t5("cuda", max_length=256 if is_schnell else 512)
@@ -109,7 +109,7 @@ def b64(txt,vec):
 
 t5,clip=load_encoders()
 
-@spaces.GPU(duration=10)
+@spaces.GPU(duration=20)
 def convert(prompt):
     if isinstance(prompt, str):
         prompt = [prompt]
@@ -118,6 +118,25 @@ def convert(prompt):
     vec = clip.tokenize(prompt)
     vec = clip(vec)
     return b64(txt,vec)
+import jax
+def _to_embed(t5, clip, txt, vec):
+    t5=nnx.merge(*t5)
+    clip=nnx.merge(*clip)
+    return t5(txt), clip(vec)
+
+to_embed=jax.jit(_to_embed)
+
+t5_tuple=nnx.split(t5)
+clip_tuple=nnx.split(clip)
+
+@spaces.GPU(duration=120)
+def compile(prompt):
+    if isinstance(prompt, str):
+        prompt = [prompt]
+    txt = t5.tokenize(prompt)
+    vec = clip.tokenize(prompt)
+    text,vec=to_embed(t5_tuple,clip_tuple,txt,vec)
+    return b64(txt,vec)
 
 with gr.Blocks() as demo:
     gr.Markdown("""A workaround for flux-flax to fit into 40G VRAM""")
@@ -125,10 +144,14 @@ with gr.Blocks() as demo:
         with gr.Column():
             prompt = gr.Textbox(label="prompt")
             convert_btn = gr.Button(value="Convert")
+            compile_btn = gr.Button(value="Compile")
         with gr.Column():
             output = gr.Textbox(label="output")
 
     convert_btn.click(convert, inputs=prompt, outputs=output, api_name="convert")
+    compile_btn.click(compile, inputs=prompt, outputs=output, api_name="compile")
+
+
 
 
 demo.launch()