Update app.py

app.py

@@ -52,6 +52,14 @@ WAN21_FILES = [
     "Wan2.1_VAE.pth"
 ]
 
+# Define tokenizer files to download
+UMT5_XXL_TOKENIZER_FILES = [
+    "google/umt5-xxl/special_tokens_map.json",
+    "google/umt5-xxl/spiece.model",
+    "google/umt5-xxl/tokenizer.json",
+    "google/umt5-xxl/tokenizer_config.json"
+]
+
 RECAMMASTER_REPO_ID = "KwaiVGI/ReCamMaster-Wan2.1"
 RECAMMASTER_CHECKPOINT_FILE = "step20000.ckpt"
 RECAMMASTER_LOCAL_DIR = f"{MODELS_ROOT_DIR}/ReCamMaster/checkpoints"
@@ -59,6 +67,55 @@ RECAMMASTER_LOCAL_DIR = f"{MODELS_ROOT_DIR}/ReCamMaster/checkpoints"
 # Define test data directory
 TEST_DATA_DIR = "example_test_data"
 
+def download_umt5_xxl_tokenizer(progress_callback=None):
+    """Download UMT5-XXL tokenizer files from HuggingFace"""
+    
+    total_files = len(UMT5_XXL_TOKENIZER_FILES)
+    downloaded_paths = []
+    
+    for i, file_path in enumerate(UMT5_XXL_TOKENIZER_FILES):
+        local_dir = f"{WAN21_LOCAL_DIR}/{os.path.dirname(file_path)}"
+        filename = os.path.basename(file_path)
+        full_local_path = f"{WAN21_LOCAL_DIR}/{file_path}"
+        
+        # Update progress
+        if progress_callback:
+            progress_callback(i/total_files, desc=f"Checking tokenizer file {i+1}/{total_files}: {filename}")
+        
+        # Check if already exists
+        if os.path.exists(full_local_path):
+            logger.info(f"✓ Tokenizer file {filename} already exists at {full_local_path}")
+            downloaded_paths.append(full_local_path)
+            continue
+        
+        # Create directory if it doesn't exist
+        os.makedirs(local_dir, exist_ok=True)
+        
+        # Download the file
+        logger.info(f"Downloading tokenizer file {filename} from {WAN21_REPO_ID}/{file_path}...")
+        
+        if progress_callback:
+            progress_callback(i/total_files, desc=f"Downloading tokenizer file {i+1}/{total_files}: {filename}")
+        
+        try:
+            # Download using huggingface_hub
+            downloaded_path = hf_hub_download(
+                repo_id=WAN21_REPO_ID,
+                filename=file_path,
+                local_dir=WAN21_LOCAL_DIR,
+                local_dir_use_symlinks=False
+            )
+            logger.info(f"✓ Successfully downloaded tokenizer file {filename} to {downloaded_path}!")
+            downloaded_paths.append(downloaded_path)
+        except Exception as e:
+            logger.error(f"✗ Error downloading tokenizer file {filename}: {e}")
+            raise
+    
+    if progress_callback:
+        progress_callback(1.0, desc=f"All tokenizer files downloaded successfully!")
+    
+    return downloaded_paths
+
 def download_wan21_models(progress_callback=None):
     """Download Wan2.1 model files from HuggingFace"""
     
@@ -343,10 +400,45 @@ def load_models(progress_callback=None):
             logger.error(error_msg)
             return error_msg
         
+        # Fourth, download UMT5-XXL tokenizer files
+        if progress_callback:
+            progress_callback(0.3, desc="Checking for UMT5-XXL tokenizer files...")
+        
+        try:
+            tokenizer_paths = download_umt5_xxl_tokenizer(progress_callback)
+            logger.info(f"Using UMT5-XXL tokenizer files: {tokenizer_paths}")
+        except Exception as e:
+            error_msg = f"Error downloading UMT5-XXL tokenizer files: {str(e)}"
+            logger.error(error_msg)
+            return error_msg
+        
         # Now, load the models
         if progress_callback:
             progress_callback(0.4, desc="Loading model manager...")
         
+        # Create symlink for google/umt5-xxl to handle potential path issues
+        # Some libraries might look for this in a different way
+        try:
+            google_dir = f"{MODELS_ROOT_DIR}/google"
+            if not os.path.exists(google_dir):
+                os.makedirs(google_dir, exist_ok=True)
+            
+            umt5_xxl_symlink = f"{google_dir}/umt5-xxl"
+            umt5_xxl_source = f"{WAN21_LOCAL_DIR}/google/umt5-xxl"
+            
+            # Create a symlink if it doesn't exist
+            if not os.path.exists(umt5_xxl_symlink) and os.path.exists(umt5_xxl_source):
+                if os.name == 'nt':  # Windows
+                    import ctypes
+                    kdll = ctypes.windll.LoadLibrary("kernel32.dll")
+                    kdll.CreateSymbolicLinkA(umt5_xxl_symlink.encode(), umt5_xxl_source.encode(), 1)
+                else:  # Unix/Linux
+                    os.symlink(umt5_xxl_source, umt5_xxl_symlink)
+                logger.info(f"Created symlink from {umt5_xxl_source} to {umt5_xxl_symlink}")
+        except Exception as e:
+            logger.warning(f"Could not create symlink for google/umt5-xxl: {str(e)}")
+            # This is a warning, not an error, as we'll try to proceed anyway
+        
         # Load Wan2.1 pre-trained models
         model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
         
@@ -363,6 +455,13 @@ def load_models(progress_callback=None):
                 logger.error(error_msg)
                 return error_msg
         
+        # Set environment variable for transformers to find the tokenizer
+        os.environ["TRANSFORMERS_CACHE"] = MODELS_ROOT_DIR
+        
+        # Set the configuration for the text encoder to use the downloaded tokenizer path
+        # This is needed because the WanTextEncoder expects the tokenizer to be at this path
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Disable tokenizers parallelism warning
+        
         model_manager.load_models(model_files)
         
         if progress_callback:
@@ -581,13 +680,23 @@ def generate_recammaster_video(
         return None, f"Error: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="ReCamMaster") as demo:
-
+with gr.Blocks(title="ReCamMaster Demo") as demo:
+    # Show loading status
+    loading_status = gr.Textbox(
+        label="Model Loading Status", 
+        value="Loading models, please wait...",
+        interactive=False,
+        visible=True
+    )
+    
     gr.Markdown(f"""
-    # 🎥 ReCamMaster
+    # 🎥 ReCamMaster Demo
     
     ReCamMaster allows you to re-capture videos with novel camera trajectories.
     Upload a video and select a camera transformation to see the magic!
+    
+    **Note:** All required models will be automatically downloaded to {MODELS_ROOT_DIR} when you start the app.
+    You can customize this location by setting the RECAMMASTER_MODELS_DIR environment variable.
     """)
     
     with gr.Row():
@@ -618,8 +727,24 @@ with gr.Blocks(title="ReCamMaster") as demo:
             # Output section
             output_video = gr.Video(label="Output Video")
             status_output = gr.Textbox(label="Generation Status", interactive=False)
-
-
+    
+    # Example videos
+    gr.Markdown("### Example Videos")
+    gr.Examples(
+        examples=[
+            [f"{TEST_DATA_DIR}/videos/case0.mp4", "A person dancing", "1"],
+            [f"{TEST_DATA_DIR}/videos/case1.mp4", "A scenic view", "5"],
+        ],
+        inputs=[video_input, text_prompt, camera_type],
+    )
+    
+    # Load models automatically when the interface loads
+    def on_load():
+        status = load_models()
+        return gr.update(value=status, visible=True if "Error" in status else False)
+    
+    demo.load(on_load, outputs=[loading_status])
+    
     # Event handlers
     generate_btn.click(
         fn=generate_recammaster_video,
@@ -628,5 +753,4 @@ with gr.Blocks(title="ReCamMaster") as demo:
     )
 
 if __name__ == "__main__":
-    load_models()
     demo.launch(share=True)