Skip to content

Commit

Permalink
add a way to specify the subfolder of the unet
Browse files Browse the repository at this point in the history
(no subfolder -> pass an empty string)
  • Loading branch information
catwell committed Jan 16, 2024
1 parent 09c358a commit 91b7423
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion scripts/conversion/convert_diffusers_unet.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
class Args(argparse.Namespace):
source_path: str
output_path: str | None
subfolder: str | None
half: bool
verbose: bool

Expand All @@ -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`
source: nn.Module = UNet2DConditionModel.from_pretrained( # type: ignore
pretrained_model_name_or_path=args.source_path,
subfolder="unet",
subfolder=args.subfolder,
low_cpu_mem_usage=False,
)
source_in_channels: int = source.config.in_channels # type: ignore
Expand Down Expand Up @@ -81,6 +82,7 @@ def main() -> None:
" 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(
"--verbose",
Expand Down

0 comments on commit 91b7423

Please sign in to comment.