DiffusionLens / diffusion_lens.py
tokeron's picture
Update diffusion_lens.py
84d321b verified
raw
history blame
1.31 kB
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch
import numpy as np
import gradio as gr
model_dict = {
'Stable Diffusion 1.4': "CompVis/stable-diffusion-v1-4",
'Stable Diffusion 2.1': "stabilityai/stable-diffusion-2-1",
}
model_num_of_layers = {
'Stable Diffusion 1.4': 12,
'Stable Diffusion 2.1': 22,
}
# global variable
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32
def get_images(prompt, skip_layers, model, seed):
model_name = model_dict[model]
pipeline = StableDiffusionPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
variant="fp16",
add_watermarker=False,
)
# Move the pipeline to the device
pipeline.to(device)
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
print('inside get images')
layer = model_num_of_layers[model] - skip_layers
gr.Info(f"Generating an image layer number {layer}")
print(f'skipping {skip_layers}')
pipeline_output = pipeline(prompt, clip_skip=skip_layers, num_images_per_prompt=1, return_tensors=False, seed=seed)
print('after pipeline')
images = pipeline_output.images
print('got images')
return images