feat(api): add option for HuggingFace token in convert script
This commit is contained in:
parent
2a6df0f3aa
commit
45a3ddc2a9
|
@ -182,7 +182,7 @@ def onnx_export(
|
||||||
|
|
||||||
|
|
||||||
@torch.no_grad()
|
@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
|
From https://github.com/huggingface/diffusers/blob/main/scripts/convert_stable_diffusion_checkpoint_to_onnx.py
|
||||||
'''
|
'''
|
||||||
|
@ -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')
|
'Half precision model export is only supported on GPUs with CUDA')
|
||||||
|
|
||||||
pipeline = StableDiffusionPipeline.from_pretrained(
|
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)
|
output_path = Path(dest_path)
|
||||||
|
|
||||||
# TEXT ENCODER
|
# TEXT ENCODER
|
||||||
|
@ -387,7 +387,7 @@ def load_models(args, models: Models):
|
||||||
if source[0] in args.skip:
|
if source[0] in args.skip:
|
||||||
print('Skipping model: %s' % source[0])
|
print('Skipping model: %s' % source[0])
|
||||||
else:
|
else:
|
||||||
convert_diffuser(*source, args.opset, args.half)
|
convert_diffuser(*source, args.opset, args.half, args.token)
|
||||||
|
|
||||||
if args.upscaling:
|
if args.upscaling:
|
||||||
for source in models.get('upscaling'):
|
for source in models.get('upscaling'):
|
||||||
|
@ -429,6 +429,11 @@ def main() -> int:
|
||||||
type=int,
|
type=int,
|
||||||
help="The version of the ONNX operator set to use.",
|
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()
|
args = parser.parse_args()
|
||||||
print(args)
|
print(args)
|
||||||
|
|
Loading…
Reference in New Issue