hack_spb / app.py
Vsevolod's picture
own token
03b568b
raw
history blame
2.21 kB
import torch
import gradio as gr
import numpy as np
from torch import autocast
from PIL import Image
from diffusers import StableDiffusionImg2ImgPipeline
# load the pipeline
device = "cpu"#"cuda"
model_id_or_path = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
model_id_or_path,
revision = "fp32", #fp16"б
torch_dtype = torch.float32, #torch.float16,
use_auth_token = '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'
)
# or download via git clone https://huggingface.co/CompVis/stable-diffusion-v1-4
# and pass `model_id_or_path="./stable-diffusion-v1-4"` without having to use `use_auth_token=True`.
pipe = pipe.to(device)
def diffuse(x, param):
print('in callback')
x = Image.fromarray(np.uint8(x))
init_image = x.resize((768, 512))
prompt = 'st petersburg logo'
if param == 'Эрмитаж':
prompt = "st petersburg logo winter palace image on background hermitage vector style"
elif param == 'Казанский собор':
prompt = "st petersburg logo kazansky sobor image on background"
elif param == 'Мосты':
prompt = 'st petersburg logo bridges over neva image on background beutiful high quality'
with autocast("cuda"):
images = pipe(prompt=prompt, init_image=init_image, strength=0.7, guidance_scale=7.5).images
return [images[0], param]
def flip_image(x, param):
return [np.fliplr(x), 'функция приняла на вход ' + param]
with gr.Blocks() as demo:
gr.Markdown("Слово 'Санкт-Петербург'")
with gr.Tab("Санкт-Петербург"):
with gr.Row():
image_input = gr.Image()
param_input = gr.Radio(["Эрмитаж", "Мосты", "Казанский собор"], label='Что для тебя Санкт-Петербург?')
image_output = gr.Image()
param_out = gr.Markdown()
image_button = gr.Button("GET IMAGE")
image_button.click(diffuse, [image_input, param_input], [image_output, param_out])
demo.launch()
#def greet(name):
# return "Hello " + name + "!!"
#iface = gr.Interface(fn=greet, inputs="text", outputs="text")
#iface.launch()