Skip to content

Instantly share code, notes, and snippets.

@lmmx
Last active August 31, 2022 18:38
Embed
What would you like to do?
Demo of Stable Diffusion usage, storing the prompt in metadata https://twitter.com/permutans/status/1562471438501548034
import piexif
from torch import autocast
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4",
use_auth_token=True,
).to("cuda")
prompt = "a photo of a lilac and tan French Bulldog puppy frolicking in a park filled with tennis balls and butterflies, DSLR camera, high quality, sunny day, joyful feeling"
prompt_short = "fb-park"
def make_prompt_exif(prompt):
"""https://stackoverflow.com/a/63649983/"""
exif_ifd = {piexif.ExifIFD.UserComment: f"Prompt: {prompt}".encode()}
exif_dict = {"0th": {}, "Exif": exif_ifd, "1st": {}, "thumbnail": None, "GPS": {}}
return piexif.dump(exif_dict)
for i in range(1):
with autocast("cuda"):
size = int(512 * 1.5)
image = pipe(
prompt,
height=size,
width=size,
num_inference_steps=100,
guidance_scale=11.0,
)["sample"][0]
image.save(f"{prompt_short}_{i:03d}.png", exif=make_prompt_exif(prompt))
from pathlib import Path
import piexif
import torch
from PIL import Image
from torch import autocast
from image_to_image import StableDiffusionImg2ImgPipeline, preprocess
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4",
revision="fp16",
torch_dtype=torch.float16,
use_auth_token=True,
).to("cuda")
prompt = "drawing of a dog, pencil sketch"
prompt_short = "draw-the-dog"
size = 512 # int(512 * 1.5)
def make_prompt_exif(prompt):
"""https://stackoverflow.com/a/63649983/"""
exif_ifd = {piexif.ExifIFD.UserComment: f"Prompt: {prompt}".encode()}
exif_dict = {"0th": {}, "Exif": exif_ifd, "1st": {}, "thumbnail": None, "GPS": {}}
return piexif.dump(exif_dict)
def load_and_preprocess_init_image(image_path):
init_image = Image.open(image_path).convert("RGB").resize((size, size))
return preprocess(init_image)
image_path = Path("sketches") / "dog-ipad-sketch.jpg"
init_image = load_and_preprocess_init_image(image_path)
for i in range(100):
with autocast("cuda"):
image = pipe(
prompt,
init_image=init_image,
num_inference_steps=100,
guidance_scale=11.0,
)["sample"][0]
image.save(f"{prompt_short}_{i:03d}.png", exif=make_prompt_exif(prompt))
import json
from pathlib import Path
import piexif
import torch
from PIL import Image
from torch import autocast
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
from diffusers.schedulers import DDIMScheduler
ddim_sched = scheduler = DDIMScheduler(
beta_start=0.00085,
beta_end=0.012,
clip_sample=False,
set_alpha_to_one=False,
)
pipe = StableDiffusionInpaintPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4",
revision="fp16",
torch_dtype=torch.float16,
use_auth_token=True,
scheduler=ddim_sched,
).to("cuda")
prompt = "photo of a man smiling"
prompt_short = "me-smiling-10"
size = 512 + int(32 * 10)
strength = 0.7
num_inference_steps = 150
guidance_scale = 11.0
def make_prompt_exif(**kwargs):
"""https://stackoverflow.com/a/63649983/"""
exif_ifd = {piexif.ExifIFD.UserComment: f"Params: {json.dumps(kwargs)}".encode()}
exif_dict = {"0th": {}, "Exif": exif_ifd, "1st": {}, "thumbnail": None, "GPS": {}}
return piexif.dump(exif_dict)
def load_image(image_path):
return Image.open(image_path).convert("RGB").resize((size, size))
image_path = Path("inputs") / "twitter-pic.jpg"
mask_path = Path("inputs") / "twitter-pic-mask.jpg"
init_image = load_image(image_path)
mask_image = load_image(mask_path)
out_dir = Path("results") / prompt_short
out_dir.mkdir(exist_ok=False)
for i in range(100):
with autocast("cuda"):
image = pipe(
prompt=prompt,
strength=strength,
init_image=init_image,
mask_image=mask_image,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
)["sample"][0]
image.save(
str(out_dir / f"{prompt_short}_{i:03d}.png"),
exif=make_prompt_exif(
prompt=prompt,
prompt_short=prompt_short,
strength=strength,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
size=size,
),
)
conda create -n diffusers
conda activate diffusers
conda install -y pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c conda-forge
pip install diffusers transformers piexif

OUTDATED For hf_stablediffusion_img2img.py download https://github.com/huggingface/diffusers/blob/main/examples/inference/image_to_image.py (edit: this has been deprecated)

For hf_stablediffusion_inpaint.py run pip install git+https://github.com/huggingface/diffusers.git (not yet been released). At time of writing the DDIM scheduler worked, the default PNDM one did not.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment