crowsonkb / style-transfer-pytorch

Neural style transfer in PyTorch.
MIT License
479 stars 115 forks source link

Allow specifying custom scales and iterations #11

Open GusRuss89 opened 3 years ago

GusRuss89 commented 3 years ago

This PR allows you to optionally set custom scales and their respective iterations.

Previously, if you specify --end-scale and --iterations via the cli, all generated scales will run for the same number of iterations (apart from the initial scale which is a separate parameter).

This isn't always desirable. Larger scales can often get away with far fewer iterations than smaller scales. With this PR you can specify scales and iterations together in an array of tuples. E.g.

args = {
    'scales': [(256,1000), (512,1000), (1024, 500), (1256, 300), (1600, 200)]
}
defaults = StyleTransfer.stylize.__kwdefaults__
st_kwargs = {k: v for k, v in args.items() if k in defaults}
st.stylize(content_img, style_imgs, **st_kwargs)

This would run the 256px scale and 512px scale for 1000 iterations, 1024px for 500 iterations, 1256px for 300 and 1600px for 200. The output quality is similar and the run time is reduced massively.

This PR doesn't include an update to the CLI to set this parameter, it's currently only available when interfacing with StyleTransfer from another python file. Setting scales and iterations the previous way is still supported, just don't include the scales argument in the call to stylize.