|
import gradio as gr |
|
import torch |
|
from PIL import Image |
|
|
|
from lambda_diffusers import StableDiffusionImageEmbedPipeline |
|
|
|
def ask(input_im, scale, n_samples, steps, seed): |
|
|
|
images_list = pipe( |
|
n_samples*[input_im], |
|
guidance_scale=scale, |
|
num_inference_steps=steps, |
|
generator=generator, |
|
) |
|
for i, image in enumerate(images_list["sample"]): |
|
if(images_list["nsfw_content_detected"][i]): |
|
safe_image = Image.open(r"unsafe.png") |
|
images.append(safe_image) |
|
else: |
|
images.append(image) |
|
|
|
def main(input_im, scale, n_samples, steps, seed): |
|
|
|
generator = torch.Generator(device=device).manual_seed(int(seed)) |
|
|
|
images = [] |
|
ask(input_im, scale, n_samples, steps, seed) |
|
ask(input_im, scale, n_samples, steps, seed) |
|
|
|
return images |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
pipe = StableDiffusionImageEmbedPipeline.from_pretrained( |
|
"lambdalabs/sd-image-variations-diffusers", |
|
revision="273115e88df42350019ef4d628265b8c29ef4af5", |
|
) |
|
pipe = pipe.to(device) |
|
|
|
inputs = [ |
|
gr.Image(), |
|
gr.Slider(0, 25, value=3, step=1, label="Guidance scale"), |
|
gr.Slider(1, 2, value=2, step=1, label="Number images"), |
|
gr.Slider(5, 50, value=25, step=5, label="Steps"), |
|
gr.Number(0, labal="Seed", precision=0) |
|
] |
|
output = gr.Gallery(label="Generated variations") |
|
output.style(grid=2) |
|
|
|
demo = gr.Interface( |
|
fn=main, |
|
title="Stable Diffusion Image Variations", |
|
inputs=inputs, |
|
outputs=output |
|
) |
|
demo.launch() |
|
|