nightfury's picture
Update app.py
b7f4263
raw
history blame
1.61 kB
import gradio as gr
import torch
import streamlit as st
from PIL import Image
import numpy as np
from io import BytesIO
from diffusers import StableDiffusionImg2ImgPipeline
device="cpu"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=st.secrets['USER_TOKEN'])
pipe.to(device)
def resize(value,img):
img = Image.open(img)
img = img.resize((value,value), Image.Resampling.LANCZOS)
return img
def infer(source_img, prompt, guide, steps, seed, Strength):
generator = torch.Generator('cpu').manual_seed(seed)
source_image = resize(512, source_img)
source_image.save('source.png')
image_list = pipe([prompt], init_image=source_image, strength=Strength, guidance_scale=guide, num_inference_steps=steps)
images = []
safe_image = Image.open(r"unsafe.png")
for i, image in enumerate(image_list["sample"]):
if(image_list["nsfw_content_detected"][i]):
images.append(safe_image)
else:
images.append(image)
return image
gr.Interface(fn=infer, inputs=[gr.Image(source="upload", type="filepath", label="Raw Image"), gr.Textbox(label = 'Prompt Input Text'),
gr.Slider(2, 15, value = 7, label = 'Guidence Scale'),
gr.Slider(10, 50, value = 25, step = 1, label = 'Number of Iterations'),
gr.Slider(
label = "Seed",
minimum = 0,
maximum = 2147483647,
step = 1,
randomize = True), gr.Slider(label='Strength', minimum = 0, maximum = 1, step = .05, value = .5)
], outputs='image').queue(max_size=10).launch(enable_queue=True)