Update app.py

app.py

@@ -44,19 +44,21 @@ titok_generator = titok_generator.to(device)
 
 
 @spaces.GPU
-def demo_infer(guidance_scale, randomize_temperature, num_sample_steps,
+def demo_infer(tokenizer,
+               generator,
+               guidance_scale, randomize_temperature, num_sample_steps,
                class_label, seed):
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    titok_tokenizer = titok_tokenizer.to(device)
-    titok_generator = titok_generator.to(device)
+    tokenizer = tokenizer.to(device)
+    generator = generator.to(device)
     n = 4
     class_labels = [class_label for _ in range(n)]
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)
     t1 = time.time()
     generated_image = demo_util.sample_fn(
-        generator=titok_generator,
-        tokenizer=titok_tokenizer,
+        generator=generator,
+        tokenizer=tokenizer,
         labels=class_labels,
         guidance_scale=guidance_scale,
         randomize_temperature=randomize_temperature,
@@ -90,6 +92,7 @@ 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])