fixed model loading bug

app.py

@@ -47,14 +47,14 @@ def get_num_transfer_tokens(mask_index, steps):
     return num_transfer_tokens
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-DEFAULT_MODEL_PATH = "Gen-Verse/MMaDA-8B-MixCoT" # Default
-MASK_ID = 126336 
-MODEL = MMadaModelLM.from_pretrained(DEFAULT_MODEL_PATH, trust_remote_code=True, torch_dtype=torch.bfloat16).to(DEVICE).eval()
-TOKENIZER = AutoTokenizer.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-MixCoT" 
+MASK_ID = None  # 初始化为 None
+MODEL = None    # 初始化为 None
+TOKENIZER = None# 初始化为 None
+uni_prompting = None # 初始化为 None
+VQ_MODEL = None # 初始化为 None, 稍后在初始化函数中加载
 
-CURRENT_MODEL_PATH = DEFAULT_MODEL_PATH
+CURRENT_MODEL_PATH = None # 初始化为 None
 
 MODEL_CHOICES = [
     "MMaDA-8B-Base",
@@ -1021,9 +1021,15 @@ with gr.Blocks(css=css_styles, theme=theme) as demo:
     )
 
     def initialize_app_state():
-        default_model_choice = "MMaDA-8B-MixCoT" # 默认加载 MixCoT
+        global VQ_MODEL 
+        
+        if VQ_MODEL is None:
+            print("Loading VQ_MODEL for the first time...")
+            VQ_MODEL = MAGVITv2().from_pretrained("showlab/magvitv2").to(DEVICE)
+            print("VQ_MODEL loaded.")
+            
+        default_model_choice = "MMaDA-8B-MixCoT"
 
-        # handle_model_selection_change 现在返回更多项
         status, lm_b_vis, lm_m_vis, lm_x_vis, \
         mmu_b_vis, mmu_m_vis, mmu_x_vis, \
         init_thinking_lm_state, init_think_lm_btn_update, \
@@ -1031,12 +1037,12 @@ with gr.Blocks(css=css_styles, theme=theme) as demo:
         
         return (
             default_model_choice, 
-            status,               
-            lm_b_vis,             
-            lm_m_vis,             
-            lm_x_vis,             
-            mmu_b_vis,            
-            mmu_m_vis,            
+            status,             
+            lm_b_vis,            
+            lm_m_vis,            
+            lm_x_vis,            
+            mmu_b_vis,           
+            mmu_m_vis,           
             mmu_x_vis,
             init_thinking_lm_state,
             init_think_lm_btn_update,
@@ -1056,10 +1062,10 @@ with gr.Blocks(css=css_styles, theme=theme) as demo:
             examples_mmu_base,
             examples_mmu_mixcot,
             examples_mmu_max,
-            thinking_mode_lm,       # gr.State for LM thinking mode
-            think_button_lm,        # gr.Button for LM thinking mode
-            thinking_mode_mmu,      # gr.State for MMU thinking mode
-            think_button_mmu        # gr.Button for MMU thinking mode
+            thinking_mode_lm,
+            think_button_lm,
+            thinking_mode_mmu,
+            think_button_mmu
         ],
         queue=True
     )