Update

app.py

@@ -12,7 +12,7 @@ from jax import Array as Tensor
 from transformers import (FlaxCLIPTextModel, CLIPTokenizer, FlaxT5EncoderModel,
                           T5Tokenizer)
 
-
+models = {}
 class HFEmbedder(nnx.Module):
     def __init__(self, version: str, max_length: int, **hf_kwargs):
         self.is_clip = version.startswith("openai")
@@ -29,7 +29,7 @@ class HFEmbedder(nnx.Module):
             self.hf_module, params = FlaxT5EncoderModel.from_pretrained(version, _do_init=False,**hf_kwargs)
         self.hf_module._is_initialized = True
         import jax
-        self.hf_module.params = jax.tree_map(lambda x: jax.device_put(x, jax.devices("cuda")[0]), params)
+        self.hf_module.params = jax.tree.map(lambda x: jax.device_put(x, jax.devices("cuda")[0]), params)
         # if dtype==jnp.bfloat16:
 
     def tokenize(self, text: list[str]) -> Tensor:
@@ -107,10 +107,11 @@ def b64(txt,vec):
     encoded = base64.b64encode(buffer.getvalue()).decode('utf-8')
     return encoded
 
-t5,clip=load_encoders()
+# t5,clip=load_encoders()
 
 @spaces.GPU(duration=20)
 def convert(prompt):
+    t5,clip=models["t5"],models["clip"]
     if isinstance(prompt, str):
         prompt = [prompt]
     txt = t5.tokenize(prompt)
@@ -126,11 +127,12 @@ def _to_embed(t5, clip, txt, vec):
 
 to_embed=jax.jit(_to_embed)
 
-t5_tuple=nnx.split(t5)
-clip_tuple=nnx.split(clip)
+# t5_tuple=nnx.split(t5)
+# clip_tuple=nnx.split(clip)
 
 @spaces.GPU(duration=120)
 def compile(prompt):
+    t5,clip,t5_tuple,clip_tuple=models["t5"],models["clip"],models["t5_tuple"],models["clip_tuple"]
     if isinstance(prompt, str):
         prompt = [prompt]
     txt = t5.tokenize(prompt)
@@ -138,6 +140,17 @@ def compile(prompt):
     text,vec=to_embed(t5_tuple,clip_tuple,txt,vec)
     return b64(txt,vec)
 
+@spaces.GPU(duration=120)
+def load(prompt):
+    is_schnell = True
+    t5 = load_t5("cuda", max_length=256 if is_schnell else 512)
+    clip = load_clip("cuda")
+    models["t5"]=t5
+    models["clip"]=clip
+    models["t5_tuple"]=nnx.split(t5)
+    models["clip_tuple"]=nnx.split(clip)
+    return "Loaded"
+
 with gr.Blocks() as demo:
     gr.Markdown("""A workaround for flux-flax to fit into 40G VRAM""")
     with gr.Row():
@@ -145,9 +158,10 @@ with gr.Blocks() as demo:
             prompt = gr.Textbox(label="prompt")
             convert_btn = gr.Button(value="Convert")
             compile_btn = gr.Button(value="Compile")
+            load_btn = gr.Button(value="Load")
         with gr.Column():
             output = gr.Textbox(label="output")
-
+    load_btn.click(load, inputs=prompt, outputs=output, api_name="load")
     convert_btn.click(convert, inputs=prompt, outputs=output, api_name="convert")
     compile_btn.click(compile, inputs=prompt, outputs=output, api_name="compile")