Created
August 28, 2022 15:04
-
-
Save Norod/b49be5918c0499a986f030a5889dda2b to your computer and use it in GitHub Desktop.
Stable Diffusion, running on CPU, uses hugging-face diffusers library
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#### pip install diffusers==0.2.4 transformers scipy ftfy | |
#### | |
from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler | |
import torch | |
def main(): | |
seed = 1000 #1000, 42, 420 | |
torch.manual_seed(seed) | |
generator = torch.Generator() | |
generator.manual_seed(seed) | |
# this will substitute the default PNDM scheduler for K-LMS | |
lms = LMSDiscreteScheduler( | |
beta_start=0.00085, | |
beta_end=0.012, | |
beta_schedule="scaled_linear", | |
tensor_format = 'pt' | |
) | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"CompVis/stable-diffusion-v1-4", | |
scheduler=lms, | |
generator=generator, | |
use_auth_token=True, | |
).to("cpu") | |
num_inference_steps = 10 | |
width=448 | |
height=512 | |
prompt = "a photo of an astronaut riding a horse on mars" | |
result = pipe(prompt, num_inference_steps=num_inference_steps, height=height, width=width) | |
nsfw_content_detected = result["nsfw_content_detected"][0] | |
print("nsfw_content_detected: " + str(nsfw_content_detected)) | |
if nsfw_content_detected: | |
prompt = "NSFW-" + prompt | |
image = result["sample"][0] | |
output_file = prompt.replace(" ", "_") + "-" + str(width) + "x" +str(height)+ "_" + str(num_inference_steps) + "steps" + "_seed" + str(seed) + ".jpg" | |
image.save(output_file) | |
print("Saved: " + str(output_file)) | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment