Update app.py

app.py

@@ -68,6 +68,7 @@ MODEL_ACTUAL_PATHS = {
 def clear_outputs_action():
         return None, None
 
+@spaces.GPU
 def _load_model_and_tokenizer_core(model_path_to_load, model_display_name_for_status):
     global MODEL, TOKENIZER, MASK_ID, CURRENT_MODEL_PATH, DEVICE, uni_prompting
     
@@ -77,42 +78,42 @@ def _load_model_and_tokenizer_core(model_path_to_load, model_display_name_for_st
     CURRENT_MODEL_PATH = model_path_to_load 
 
     status_msg_parts = [f"Loading '{model_display_name_for_status}'..."]
-    try:
-        TOKENIZER = AutoTokenizer.from_pretrained(model_path_to_load, trust_remote_code=True)
-        status_msg_parts.append(f"Tokenizer for '{model_display_name_for_status}' loaded.")
+    # try:
+    TOKENIZER = AutoTokenizer.from_pretrained(model_path_to_load, trust_remote_code=True)
+    status_msg_parts.append(f"Tokenizer for '{model_display_name_for_status}' loaded.")
 
-        MODEL = MMadaModelLM.from_pretrained(model_path_to_load, trust_remote_code=True, torch_dtype=torch.bfloat16).to(DEVICE).eval()
-        status_msg_parts.append(f"Model '{model_display_name_for_status}' loaded to {DEVICE}.")
+    MODEL = MMadaModelLM.from_pretrained(model_path_to_load, trust_remote_code=True, torch_dtype=torch.bfloat16).to(DEVICE).eval()
+    status_msg_parts.append(f"Model '{model_display_name_for_status}' loaded to {DEVICE}.")
 
-        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)
-        
-        if hasattr(TOKENIZER, 'mask_token_id') and TOKENIZER.mask_token_id is not None:
-            MASK_ID = TOKENIZER.mask_token_id
-            status_msg_parts.append(f"Using MASK_ID from tokenizer: {MASK_ID}.")
+    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)
+    
+    if hasattr(TOKENIZER, 'mask_token_id') and TOKENIZER.mask_token_id is not None:
+        MASK_ID = TOKENIZER.mask_token_id
+        status_msg_parts.append(f"Using MASK_ID from tokenizer: {MASK_ID}.")
+    else:
+        MASK_ID = 126336 
+        status_msg_parts.append(f"Using default MASK_ID: {MASK_ID}.")
+
+    if TOKENIZER.pad_token_id is None:
+        if TOKENIZER.eos_token_id is not None:
+            TOKENIZER.pad_token_id = TOKENIZER.eos_token_id
+            TOKENIZER.pad_token = TOKENIZER.eos_token 
+            status_msg_parts.append(f"Set pad_token_id to eos_token_id ({TOKENIZER.eos_token_id}).")
         else:
-            MASK_ID = 126336 
-            status_msg_parts.append(f"Using default MASK_ID: {MASK_ID}.")
-
-        if TOKENIZER.pad_token_id is None:
-            if TOKENIZER.eos_token_id is not None:
-                TOKENIZER.pad_token_id = TOKENIZER.eos_token_id
-                TOKENIZER.pad_token = TOKENIZER.eos_token 
-                status_msg_parts.append(f"Set pad_token_id to eos_token_id ({TOKENIZER.eos_token_id}).")
-            else:
-                status_msg_parts.append("Warning: pad_token_id is None and no eos_token_id.")
-        
-        if TOKENIZER.eos_token_id is None: # Important for cleaning up output in visualization
-             status_msg_parts.append("Warning: tokenizer.eos_token_id is None. EOS cleanup might not work.")
+            status_msg_parts.append("Warning: pad_token_id is None and no eos_token_id.")
+    
+    if TOKENIZER.eos_token_id is None: # Important for cleaning up output in visualization
+         status_msg_parts.append("Warning: tokenizer.eos_token_id is None. EOS cleanup might not work.")
 
-        TOKENIZER.chat_template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n' }}"
-        
-        return " ".join(status_msg_parts)
-    except Exception as e:
-        MODEL = None
-        TOKENIZER = None
-        MASK_ID = None
-        CURRENT_MODEL_PATH = None 
-        return f"Error loading model '{model_display_name_for_status}': {str(e)}"
+    TOKENIZER.chat_template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n' }}"
+    
+    return " ".join(status_msg_parts)
+    # except Exception as e:
+    #     MODEL = None
+    #     TOKENIZER = None
+    #     MASK_ID = None
+    #     CURRENT_MODEL_PATH = None 
+    #     return f"Error loading model '{model_display_name_for_status}': {str(e)}"
 
 def handle_model_selection_change(selected_model_name_ui):
     if "coming soon" in selected_model_name_ui.lower():