Skip to content

Instantly share code, notes, and snippets.

@cgcardona
Last active September 25, 2023 02:04
Show Gist options
  • Save cgcardona/26fa430a6d9d879105528c13b450c733 to your computer and use it in GitHub Desktop.
Save cgcardona/26fa430a6d9d879105528c13b450c733 to your computer and use it in GitHub Desktop.
# all the imports
import torch, time, math, os
# import specified modules
from torch import autocast
from diffusers import StableDiffusionPipeline
# confirm GPU supports the NVIDIA machine learning toolkit
assert torch.cuda.is_available()
# Stable Diffusion v1.4: CompVis/stable-diffusion-v1-4
# Stable Diffusion v1.5: runwayml/stable-diffusion-v1-5
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_auth_token=True).to("cuda")
# the prompt which will be used to create the image
prompt = "Masterpiece, Stunning, Best quality, a captivating space-themed artwork featuring the ethereal beauty of swirling nebulas and twinkling stars with an astronaut in it dark, blue themed"
# get first 20 chars of the prompt to use as the file name
first_20_chars = prompt[0:20]
# sanitize the file name by removing any "," characters
sanitized_chars = first_20_chars.replace(",", "")
sanitized_chars = sanitized_chars.replace(" ", "-")
with autocast("cuda"):
image = pipe(prompt).images[0]
# save newly generated image assets to the generate-assets/ directory
timestamp = math.ceil(time.time())
title = f"{sanitized_chars}-{timestamp}"
parent_dir = "/path/to/generated-assets"
# each time inference.py is run a new directory should be created which is named
# the current timestamp. This new directory is where the newly generated
# image should be saved
directory = title
file_path = os.path.join(parent_dir, directory)
# set permissions
mode = 0o744
os.mkdir(file_path, mode)
file_name = f"{title}.png"
file_path_and_name = f"{file_path}/{file_name}"
image.save(file_path_and_name)
print(f"{file_path_and_name} created!")
# success
# ✨ 😎 ✨
sparkle = "\U00002728"
sunglasses = "\U0001F60E"
print (f"{sparkle} {sunglasses} {sparkle}")
print ("Winning")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment