Update app.py

app.py

@@ -156,6 +156,7 @@ def get_highlighted_text_tuples(current_x_ids_batch, prompt_input_ids, prompt_le
     return intermediate_tuples
 
 @torch.no_grad()
+@spaces.GPU
 def generate_viz_wrapper_t2i(prompt_text, steps, guidance_scale, mask_schedule="cosine"):
     global MODEL, TOKENIZER, MASK_ID, DEVICE, uni_prompting
 
@@ -196,6 +197,7 @@ def generate_viz_wrapper_t2i(prompt_text, steps, guidance_scale, mask_schedule="
         
 
 @torch.no_grad()
+@spaces.GPU
 def generate_viz_wrapper_lm(prompt_text, steps, gen_length, block_length, temperature,
                          cfg_scale, remasking_strategy, thinking_mode_lm=False): 
     global MODEL, TOKENIZER, MASK_ID, DEVICE
@@ -349,6 +351,7 @@ def generate_viz_wrapper_lm(prompt_text, steps, gen_length, block_length, temper
     yield get_highlighted_text_tuples(x, input_ids, prompt_len, TOKENIZER, MASK_ID, raw_prompt_attention_mask), final_text_str
 
 @torch.no_grad()
+@spaces.GPU
 def generate_viz_wrapper(uploaded_image_pil, prompt_text, steps, gen_length, block_length, temperature,
                          cfg_scale, remasking_strategy, thinking_mode_mmu=False): 
     global MODEL, TOKENIZER, MASK_ID, DEVICE