Stable Diffusion 🎨

#@title Install packages
!pip install diffusers==0.3.0
!pip install transformers scipy ftfy
!pip install "ipywidgets>=7,<8"

Now you can login with your user token.

#@title Hugging face login
# Enable custom widgets
from google.colab import output
output.enable_custom_widget_manager()
# Login
from huggingface_hub import notebook_login

notebook_login()
Login successful
Your token has been saved to /root/.huggingface/token
#@title Import model
import torch
from diffusers import StableDiffusionPipeline

# make sure you're logged in with `huggingface-cli login`
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=True)  
#@title Put to GPU
pipe = pipe.to("cuda")
#@title Prompt
PROMPT = "An astronaut riding a horse, black and white, photorealist" #@param [""] {allow-input: true}
FILENAME = "horse_astronaut" #@param ["horse", "frog"] {allow-input: true}
#@title Run Stable Diffusion
from torch import autocast

prompt = PROMPT if len(PROMPT) > 0 else "An astronaut riding a horse"
with autocast("cuda"):
  image = pipe(prompt).images[0]  # image here is in [PIL format](https://pillow.readthedocs.io/en/stable/)

# or if you're in a google colab you can directly display it with 
image

Save image to Drive

#@title Load dependencies
!apt-get install exiftool

from google.colab import drive
drive.mount('/content/gdrive')

import os
import datetime

def filename(FILENAME, FOLDER):
  ext = ".png"
  ls = sorted(os.listdir(FOLDER))
  counter = sum([ 1 for item in ls if FILENAME in item])
  suffix = "_" + str(counter + 1) if counter > 0 else ""
  name = FILENAME + suffix + ext
  return name
  
def metadata(imagepath):
  AUTHOR = "Samuel ORTION"
  COPYRIGHT = "Creative Commons By SA 4.0 or later"
  RELATION = "https://samuel.ortion.fr/"
  DESCRIPTION = f"prompt: {PROMPT}"
  cmd = f"""exiftool -use MWG -P -m -overwrite_original_in_place -Copyright="{COPYRIGHT}" -Creator="{AUTHOR}" -Relation="{RELATION}" -Exif:ImageDescription="{DESCRIPTION}" -Description="{DESCRIPTION}" "{imagepath}" """
  os.system(cmd)

def save(image):
  OUTFOLDER = "/content/gdrive/MyDrive/media/stablediff"
  FILEPATH = os.path.join(OUTFOLDER, filename(FILENAME, OUTFOLDER))
  image.save(FILEPATH)
  metadata(FILEPATH)
#@title Save
save(image)