@@ -1725,16 +1725,24 @@ def _add_install_flags_and_urls(self, install_cmd: list, install_flags: list, se
17251725 """Add installation flags and index URLs to command"""
17261726 # Check if this is a torch installation step by looking at the command
17271727 # Only add torch index URL if ALL packages are torch-related (torch, torchvision, torchaudio)
1728- torch_packages = ['torch' , 'torchvision' , 'torchaudio' ]
1728+ torch_packages = ['torch' , 'torchvision' , 'torchaudio' , 'xformers' ]
17291729 package_args = [arg for arg in install_cmd if not arg .startswith ('-' ) and 'pip' not in arg and 'install' not in arg and 'python' not in arg and 'uv' not in arg ]
17301730
1731- # Extract package names from version specifications like "torch>=2.4.0"
1731+ # Extract package names from version specifications like "torch>=2.4.0" and remove versions for torch packages
17321732 package_names = []
1733- for arg in package_args :
1733+ for i , arg in enumerate ( package_args ) :
17341734 # Split by common version operators
17351735 for op in ['>=' , '<=' , '==' , '!=' , '>' , '<' , '~=' ]:
17361736 if op in arg :
1737- package_names .append (arg .split (op )[0 ].strip ())
1737+ pkg_name = arg .split (op )[0 ].strip ()
1738+ package_names .append (pkg_name )
1739+ # Remove version specification for torch packages
1740+ if pkg_name .lower () in torch_packages :
1741+ # Find the index of this package in install_cmd and replace it
1742+ for j , cmd_arg in enumerate (install_cmd ):
1743+ if cmd_arg == arg :
1744+ install_cmd [j ] = pkg_name
1745+ break
17381746 break
17391747 else :
17401748 package_names .append (arg .strip ())
@@ -1970,6 +1978,9 @@ def _generate_startup_script(self, repo_path: Path, repo_info: Dict):
19701978set TMP=%tmp_path%\Temp
19711979set APPDATA=%tmp_path%\AppData\Roaming
19721980set LOCALAPPDATA=%tmp_path%\AppData\Local
1981+ set HF_HOME=%tmp_path%\huggingface
1982+ set XDG_CACHE_HOME=%tmp_path%
1983+ set HF_DATASETS_CACHE=%HF_HOME%\datasets
19731984
19741985REM Security and compatibility settings
19751986set PYTHONIOENCODING=utf-8
0 commit comments