fffiloni's picture
Update app.py
81b1f07
raw
history blame
1.55 kB
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()