add a way to specify the subfolder of the unet

(no subfolder -> pass an empty string)
This commit is contained in:
Pierre Chapuis 2024-01-16 15:16:33 +01:00
parent 94a918a474
commit d2f38871fd

View file

@ -12,6 +12,7 @@ from refiners.foundationals.latent_diffusion import SD1UNet, SDXLUNet
class Args(argparse.Namespace): class Args(argparse.Namespace):
source_path: str source_path: str
output_path: str | None output_path: str | None
subfolder: str
half: bool half: bool
verbose: bool verbose: bool
@ -20,7 +21,7 @@ def setup_converter(args: Args) -> ModelConverter:
# low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate` # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
source: nn.Module = UNet2DConditionModel.from_pretrained( # type: ignore source: nn.Module = UNet2DConditionModel.from_pretrained( # type: ignore
pretrained_model_name_or_path=args.source_path, pretrained_model_name_or_path=args.source_path,
subfolder="unet", subfolder=args.subfolder,
low_cpu_mem_usage=False, low_cpu_mem_usage=False,
) )
source_in_channels: int = source.config.in_channels # type: ignore source_in_channels: int = source.config.in_channels # type: ignore
@ -81,6 +82,7 @@ def main() -> None:
" source path." " source path."
), ),
) )
parser.add_argument("--subfolder", type=str, default="unet", help="Subfolder. Default: unet.")
parser.add_argument("--half", action="store_true", help="Convert to half precision.") parser.add_argument("--half", action="store_true", help="Convert to half precision.")
parser.add_argument( parser.add_argument(
"--verbose", "--verbose",