Skip to content

Commit

Permalink
feat(api): add option for HuggingFace token in convert script
Browse files Browse the repository at this point in the history
  • Loading branch information
ssube committed Jan 21, 2023
1 parent 2a6df0f commit 45a3ddc
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions api/onnx_web/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ def onnx_export(


@torch.no_grad()
def convert_diffuser(name: str, url: str, opset: int, half: bool):
def convert_diffuser(name: str, url: str, opset: int, half: bool, token: str):
'''
From https://github.com/huggingface/diffusers/blob/main/scripts/convert_stable_diffusion_checkpoint_to_onnx.py
'''
Expand All @@ -201,7 +201,7 @@ def convert_diffuser(name: str, url: str, opset: int, half: bool):
'Half precision model export is only supported on GPUs with CUDA')

pipeline = StableDiffusionPipeline.from_pretrained(
url, torch_dtype=dtype).to(training_device)
url, torch_dtype=dtype, use_auth_token=token).to(training_device)
output_path = Path(dest_path)

# TEXT ENCODER
Expand Down Expand Up @@ -387,7 +387,7 @@ def load_models(args, models: Models):
if source[0] in args.skip:
print('Skipping model: %s' % source[0])
else:
convert_diffuser(*source, args.opset, args.half)
convert_diffuser(*source, args.opset, args.half, args.token)

if args.upscaling:
for source in models.get('upscaling'):
Expand Down Expand Up @@ -429,6 +429,11 @@ def main() -> int:
type=int,
help="The version of the ONNX operator set to use.",
)
parser.add_argument(
'--token',
type=str,
help="HuggingFace token with read permissions for downloading models.",
)

args = parser.parse_args()
print(args)
Expand Down

0 comments on commit 45a3ddc

Please sign in to comment.