Enhance model download from HF Hub

app.py

@@ -272,6 +272,15 @@ def ui(GPU_memory_mode, scheduler_dict, config_path, compile_dit, weight_dtype):
         with gr.Column(variant="panel"):
             # Hide model selection
             diffusion_transformer_dropdown, _ = create_model_checkpoints(controller, visible=False, default_model="Wan-AI/Wan2.1-T2V-14B")
+            
+            # Use snapshot download for the VideoCoF repo to get all weights (including safetensors)
+            try:
+                from huggingface_hub import snapshot_download
+                print("Downloading VideoCoF weights...")
+                snapshot_download(repo_id="XiangpengYang/VideoCoF")
+            except Exception as e:
+                print(f"Warning: Failed to pre-download VideoCoF weights: {e}")
+
             base_model_dropdown, lora_model_dropdown, lora_alpha_slider, _ = create_finetune_models_checkpoints(controller, visible=False, default_lora="XiangpengYang/VideoCoF")
             
             # Set default LoRA alpha to 1.0 (matching inference.py)

videox_fun/models/wan_transformer3d.py

@@ -1161,6 +1161,35 @@ class WanTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     ):
         if subfolder is not None:
             pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+        
+        # Handle HF hub download for model weights
+        import os
+        from huggingface_hub import hf_hub_download, snapshot_download
+        
+        # Check if pretrained_model_path is a Hugging Face repo ID
+        if not os.path.exists(pretrained_model_path):
+            try:
+                print(f"Downloading model from HF repo: {pretrained_model_path}...")
+                # Download config.json
+                config_path = hf_hub_download(repo_id=pretrained_model_path, filename="config.json")
+                pretrained_model_path = os.path.dirname(config_path)
+                
+                # We also need model weights. 
+                # Try to download safetensors first
+                try:
+                    hf_hub_download(repo_id=pretrained_model_path, filename="diffusion_pytorch_model.safetensors")
+                except:
+                    pass
+                    
+                # Try bin
+                try:
+                    hf_hub_download(repo_id=pretrained_model_path, filename="diffusion_pytorch_model.bin")
+                except:
+                    pass
+                
+            except Exception as e:
+                print(f"Failed to download model from HF: {e}")
+
         print(f"loaded 3D transformer's pretrained weights from {pretrained_model_path} ...")
 
         config_file = os.path.join(pretrained_model_path, 'config.json')