Update app.py

app.py

@@ -10,9 +10,7 @@ import argparse
 import demo_util
 import os
 import spaces
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
+from functools import partial
 
 model2ckpt = {
     "TiTok-L-32": ("tokenizer_titok_l32.bin", "generator_titok_l32.bin"),
@@ -32,19 +30,24 @@ parser.add_argument("--seed", type=int, default=42)
 parser.add_argument("--temperature", type=float, default=1.0, help="temperature value to sample with")
 args = parser.parse_args()
 
-config = demo_util.get_config("configs/titok_l32.yaml")
-print(config)
-titok_tokenizer = demo_util.get_titok_tokenizer(config)
-print(titok_tokenizer)
-titok_generator = demo_util.get_titok_generator(config)
-print(titok_generator)
 
-titok_tokenizer = titok_tokenizer.to(device)
-titok_generator = titok_generator.to(device)
+@spaces.GPU
+def load_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    config = demo_util.get_config("configs/titok_l32.yaml")
+    print(config)
+    titok_tokenizer = demo_util.get_titok_tokenizer(config)
+    print(titok_tokenizer)
+    titok_generator = demo_util.get_titok_generator(config)
+    print(titok_generator)
+
+    titok_tokenizer = titok_tokenizer.to(device)
+    titok_generator = titok_generator.to(device)
+    return titok_tokenizer, titok_generator
 
 
 @spaces.GPU
-def demo_infer(tokenizer,
+def demo_infer_(tokenizer,
                generator,
                guidance_scale, randomize_temperature, num_sample_steps,
                class_label, seed):
@@ -70,7 +73,8 @@ def demo_infer(tokenizer,
     samples = [Image.fromarray(sample) for sample in generated_image]
     return samples
 
-
+titok_tokenizer, titok_generator = load_model()
+demo_infer = partial(demo_infer_, tokenizer=titok_tokenizer, generator=titok_generator)
 with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align: center'>An Image is Worth 32 Tokens for Reconstruction and Generation</h1>")
 
@@ -92,7 +96,6 @@ with gr.Blocks() as demo:
                 with gr.Column():
                     output = gr.Gallery(label='Generated Images', height=700)
                     button.click(demo_infer, inputs=[
-                        titok_tokenizer, titok_generator,
                         guidance_scale, randomize_temperature, num_sample_steps,
                         i1k_class, seed],
                         outputs=[output])