@@ -469,27 +469,29 @@ def validate_model_name(name, base_dir=None):
469469# relative path and nothing else
470470WEIGHT_NAME_KEY = "weight_name"
471471
472- # Tasks that read their weights with safetensors and nothing else. A pickle
473- # format is refused by name rather than left to fail in the loader: the
474- # refusal is the documented contract, not an accident of which loader runs
475- SAFETENSORS_ONLY_COMMANDS = ()
472+ # Tasks that read their weights with safetensors and nothing else, from a file
473+ # named bare (no subfolder: the task knows its default repo's layout). A
474+ # pickle format is refused by name rather than left to fail in the loader:
475+ # the refusal is the documented contract, not an accident of which loader runs
476+ SAFETENSORS_ONLY_COMMANDS = ("upscale_h3_latents" ,)
476477SAFETENSORS_SUFFIX = ".safetensors"
477478
478479
479- def validate_weight_name (name , suffixes = None , what = "weight_name" ):
480+ def validate_weight_name (name , suffixes = None , what = "weight_name" , subfolders = True ):
480481 """A file name inside a Hub repo: relative, and never climbing out of it.
481482
482483 `hf_hub_download(filename=...)` joins the name onto the local cache
483484 directory, so a name shaped like a path is a path on this machine. A
484- subfolder is legitimate (the H3 upscaler's default lives under one);
485- an absolute path, a backslash, a drive, an empty segment or a '.'/'..'
485+ subfolder is legitimate where `subfolders` allows one (an IP adapter's
486+ lives under one); an absolute path, a backslash, a drive, an empty segment or a '.'/'..'
486487 segment is not. The refusal names only what the caller wrote, never a
487488 directory on the server.
488489
489490 Args:
490491 name: The file name the workflow supplied
491492 suffixes: The file endings allowed, or None for any
492493 what: Short phrase naming the argument, for the error message
494+ subfolders: Whether the name may hold a '/' at all
493495
494496 Raises:
495497 PathTraversalError: On a name that is not a plain relative path
@@ -510,6 +512,10 @@ def validate_weight_name(name, suffixes=None, what="weight_name"):
510512 f"Refusing a { what } { name }
511513 f"path segment, so it does not name a file inside the model repo."
512514 )
515+ if not subfolders and "/" in name :
516+ raise PathTraversalError (
517+ f"Refusing a { what } { name }
518+ )
513519 if suffixes and not name .lower ().endswith (tuple (suffixes )):
514520 raise InvalidInputError (
515521 f"Refusing a { what } { name }
@@ -583,31 +589,40 @@ def location_errors(definition, source_indices=None, base_dir=None):
583589 if not isinstance (step , dict ):
584590 continue
585591 source = references .author_index (source_indices , index )
586- _walk (step , f"steps[{ source } , base_dir , errors , _weight_suffixes (step ))
592+ _walk (step , f"steps[{ source } , base_dir , errors , _weight_rules (step ))
587593 return errors
588594
589595
590- def _weight_suffixes (step ):
591- """The weight file endings a step's task reads, or None for any."""
596+ def _weight_rules (step ):
597+ """validate_weight_name's (suffixes, subfolders) for a step's task, and
598+ whether its model_name must be a Hub repo id (a task that downloads its
599+ weights reads no local directory, so one is refused here, not at run)."""
592600 task = step .get ("task" )
593601 command = task .get ("command" ) if isinstance (task , dict ) else None
594602 if command in SAFETENSORS_ONLY_COMMANDS :
595- return (SAFETENSORS_SUFFIX ,)
596- return None
603+ return (SAFETENSORS_SUFFIX ,), False , True
604+ return None , True , False
597605
598606
599- def _walk (node , path , base_dir , errors , weight_suffixes = None ):
607+ def _walk (node , path , base_dir , errors , weight_rules = ( None , True , False ) ):
600608 if isinstance (node , dict ):
601609 for key , value in node .items ():
602610 here = f"{ path } { key }
603611 if key == "model_name" and isinstance (value , str ):
604- message = _model_name_message (value , base_dir )
612+ message = _model_name_message (value , base_dir , weight_rules [ 2 ] )
605613 if message :
606614 errors .append ({"path" : here , "message" : message })
607615 continue
608616 if key == WEIGHT_NAME_KEY and isinstance (value , str ):
609617 if not _deferred (value ):
610- message = _refusal (validate_weight_name , value , weight_suffixes )
618+ suffixes , subfolders , _ = weight_rules
619+ message = _refusal (
620+ validate_weight_name ,
621+ value ,
622+ suffixes ,
623+ WEIGHT_NAME_KEY ,
624+ subfolders ,
625+ )
611626 if message :
612627 errors .append ({"path" : here , "message" : message })
613628 continue
@@ -642,10 +657,10 @@ def _walk(node, path, base_dir, errors, weight_suffixes=None):
642657 message = _check (item , base_dir , f"'{ key } )
643658 if message :
644659 errors .append ({"path" : sub_path , "message" : message })
645- _walk (value , here , base_dir , errors , weight_suffixes )
660+ _walk (value , here , base_dir , errors , weight_rules )
646661 elif isinstance (node , list ):
647662 for index , item in enumerate (node ):
648- _walk (item , f"{ path } { index } , base_dir , errors , weight_suffixes )
663+ _walk (item , f"{ path } { index } , base_dir , errors , weight_rules )
649664
650665
651666def _each (value , path ):
@@ -660,12 +675,32 @@ def _glob_message(pattern, base_dir):
660675 return _refusal (validate_media_glob , pattern , base_dir )
661676
662677
663- def _model_name_message (name , base_dir ):
678+ def _model_name_message (name , base_dir , hub_only = False ):
664679 if _deferred (name ):
665680 return None
681+ if hub_only :
682+ return _refusal (validate_hub_repo_id , name )
666683 return _refusal (validate_model_name , name , base_dir )
667684
668685
686+ def validate_hub_repo_id (name , what = "model_name" ):
687+ """A Hugging Face repo id and nothing else - for a task that downloads.
688+
689+ Raises:
690+ InvalidInputError: On a path, a URL, or any other non-repo-id
691+ """
692+ from huggingface_hub .utils import HFValidationError , validate_repo_id
693+
694+ try :
695+ validate_repo_id (str (name ))
696+ except HFValidationError :
697+ raise InvalidInputError (
698+ f"Refusing a { what } { name }
699+ f"so it takes a Hugging Face repo id ('owner/name'), not a path."
700+ ) from None
701+ return str (name )
702+
703+
669704def _refusal (check , * args ):
670705 """The message `check` refused its argument with, or None if it allowed it."""
671706 try :
0 commit comments