mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-11-24 23:28:45 +00:00
add a way to specify the subfolder of the unet
(no subfolder -> pass an empty string)
This commit is contained in:
parent
94a918a474
commit
d2f38871fd
|
@ -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",
|
||||||
|
|
Loading…
Reference in a new issue