Update app.py

app.py

@@ -45,15 +45,15 @@ def load_model():
     titok_generator = titok_generator.to(device)
     return titok_tokenizer, titok_generator
 
+titok_tokenizer, titok_generator = load_model()
 
 @spaces.GPU
-def demo_infer_(tokenizer,
-               generator,
+def demo_infer(
                guidance_scale, randomize_temperature, num_sample_steps,
                class_label, seed):
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    tokenizer = tokenizer.to(device)
-    generator = generator.to(device)
+    tokenizer = titok_tokenizer.to(device)
+    generator = titok_generator.to(device)
     n = 4
     class_labels = [class_label for _ in range(n)]
     torch.manual_seed(seed)
@@ -73,8 +73,6 @@ 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>")