divamgupta / stable-diffusion-tensorflow

Stable Diffusion in TensorFlow / Keras
Other
1.57k stars 227 forks source link

support input_image as an array, fixing minor bug in the 'load_weights_from_pytorch_ckpt' method and adding colab demo showing it's usage #50

Closed costiash closed 1 year ago

costiash commented 1 year ago

There is currently no support in "stable diffusion.py" for the case where "input_image" to the generator is not a path to an image, i.e., of type "str." The value of "input_image_tensor" will always remain "None" and the value of "input_image" will be disregarded if the value of "input_image" is an array. Added minor fix for this case. Also, the default way of pytorch to load weights is to move them to the device they were saved from. so for making the method to work we need to specify to load the weights to the cpu or changing the line: w = pt_weights['state_dict'][key].numpy() to w = pt_weights['state_dict'][key].cpu().numpy()

at least that was the case for me when trying to use this method on colab.

costiash commented 1 year ago

Added a Colab Demo to the README showing how to load new pytorch ckpt weights into the generator directly from Hugging Face.