Update app.py

app.py

@@ -46,14 +46,14 @@ def get_num_transfer_tokens(mask_index, steps):
              num_transfer_tokens[i, :remainder[i].item()] += 1 # .item() for single value tensor to int
     return num_transfer_tokens
 
-MODEL = None
-TOKENIZER = None
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-MASK_ID = None 
-uni_prompting = None
+DEFAULT_MODEL_PATH = "Gen-Verse/MMaDA-8B-Base" # Default
+MASK_ID = 126336 
+MODEL = MMadaModelLM.from_pretrained(DEFAULT_MODEL_PATH, trust_remote_code=True, torch_dtype=torch.bfloat16).to(DEVICE).eval()
+TOKENIZER = DEFAULT_MODEL_PATH.from_pretrained(DEFAULT_MODEL_PATH, trust_remote_code=True)
+uni_prompting = UniversalPrompting(TOKENIZER, max_text_len=512, special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"),ignore_id=-100, cond_dropout_prob=0.1, use_reserved_token=True)
 VQ_MODEL = MAGVITv2().from_pretrained("showlab/magvitv2").to(DEVICE)
 
-DEFAULT_MODEL_PATH = "Gen-Verse/MMaDA-8B-Base" # Default
 CURRENT_MODEL_PATH = None
 
 MODEL_CHOICES = [