Skip to content

Instantly share code, notes, and snippets.

@Norod
Created August 28, 2022 15:04
Show Gist options
  • Star 5 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Norod/b49be5918c0499a986f030a5889dda2b to your computer and use it in GitHub Desktop.
Save Norod/b49be5918c0499a986f030a5889dda2b to your computer and use it in GitHub Desktop.
Stable Diffusion, running on CPU, uses hugging-face diffusers library
#### 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