import random
import gradio as gr
from datasets import load_dataset
from PIL import Image
from set import ExpiringMap
# from model import get_sd_small, get_sd_tiny, get_sd_every
from trans_google import google_translator
import replicate
from i18n import i18nTranslator
word_list_dataset = load_dataset("Gustavosta/Stable-Diffusion-Prompts")
word_list = word_list_dataset["train"]['Prompt']
#
# from diffusers import EulerDiscreteScheduler, DDIMScheduler, KDPM2AncestralDiscreteScheduler, \
# UniPCMultistepScheduler, DPMSolverSinglestepScheduler, DEISMultistepScheduler, PNDMScheduler, \
# DPMSolverMultistepScheduler, HeunDiscreteScheduler, EulerAncestralDiscreteScheduler, DDPMScheduler, \
# LMSDiscreteScheduler, KDPM2DiscreteScheduler
# import torch
# import base64
# from io import BytesIO
is_gpu_busy = False
# translator = i18nTranslator()
# translator.init(path='locales')
samplers = [
"EulerDiscrete",
"EulerAncestralDiscrete",
"UniPCMultistep",
"DPMSolverSinglestep",
"DPMSolverMultistep",
"KDPM2Discrete",
"KDPM2AncestralDiscrete",
"DEISMultistep",
"HeunDiscrete",
"PNDM",
"DDPM",
"DDIM",
"LMSDiscrete",
]
re_sampler = [
"DDIM",
"K_EULER",
"DPMSolverMultistep",
"K_EULER_ANCESTRAL",
"PNDM",
"KLMS"
]
rand = random.Random()
translator = google_translator()
# tiny_pipe = get_sd_tiny()
# small_pipe = get_sd_small()
# every_pipe = get_sd_every()
# def get_pipe(width: int, height: int):
# if width == 512 and height == 512:
# return tiny_pipe
# elif width == 256 and height == 256:
# return small_pipe
# else:
# return every_pipe
time_client_map = ExpiringMap()
count_client_map = ExpiringMap()
def infer(prompt: str, negative: str, width: int, height: int, sampler: str,
steps: int, seed: int, scale, request: gr.Request):
client_ip = request.client.host
if client_ip != '127.0.0.1' and client_ip != 'localhost' and client_ip != '0.0.0.0':
if time_client_map.get(client_ip):
return "Too many requests from this IP, please try again later."
else:
time_client_map.put(client_ip, 1, 10) # 添加一个过期时间为 10 秒的项
count = count_client_map.get(client_ip)
if count is None:
count = 0
count += 1
if count > 5:
return "Too many requests from this IP, please try again later."
else:
count_client_map.put(client_ip, count, 24 * 60 * 60) # 添加一个过期时间为 24 小时的项
global is_gpu_busy
if seed == 0:
seed = rand.randint(0, 10000)
else:
seed = int(seed)
#
# pipeline = get_pipe(width, height)
#
images = []
# if torch.cuda.is_available():
# generator = torch.Generator(device="cuda").manual_seed(seed)
# else:
# generator = None
# if sampler == "EulerDiscrete":
# pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "EulerAncestralDiscrete":
# pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "KDPM2Discrete":
# pipeline.scheduler = KDPM2DiscreteScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "KDPM2AncestralDiscrete":
# pipeline.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "UniPCMultistep":
# pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "DPMSolverSinglestep":
# pipeline.scheduler = DPMSolverSinglestepScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "DPMSolverMultistep":
# pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "HeunDiscrete":
# pipeline.scheduler = HeunDiscreteScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "DEISMultistep":
# pipeline.scheduler = DEISMultistepScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "PNDM":
# pipeline.scheduler = PNDMScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "DDPM":
# pipeline.scheduler = DDPMScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "DDIM":
# pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
# elif sampler == "LMSDiscrete":
# pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
try:
translate_prompt = translator.translate(prompt, lang_tgt='en')
translate_negative = translator.translate(negative, lang_tgt='en')
except Exception as ex:
print(ex)
translate_prompt = prompt
translate_negative = negative
output = replicate.run(
"stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf",
input={
"prompt": translate_prompt,
"negative_prompt": translate_negative,
"guidance_scale": scale,
"num_inference_steps": steps,
"seed": seed,
"scheduler": sampler,
}
)
# image = pipeline(prompt=translate_prompt,
# negative_prompt=translate_negative,
# guidance_scale=scale,
# num_inference_steps=steps,
# generator=generator,
# height=height,
# width=width).images[0]
# buffered = BytesIO()
# image.save(buffered, format="JPEG")
# img_str = base64.b64encode(buffered.getvalue())
# img_base64 = bytes("data:image/jpeg;base64,", encoding='utf-8') + img_str
images.append(output[0])
return images
css = """
.gradio-container {
font-family: 'IBM Plex Sans', sans-serif;
}
.gr-button {
color: white;
border-color: black;
background: black;
}
input[type='range'] {
accent-color: black;
}
.dark input[type='range'] {
accent-color: #dfdfdf;
}
.container {
max-width: 1130px;
margin: auto;
padding-top: 1.5rem;
}
#prompt-column {
min-height: 500px
}
#gallery {
min-height: 22rem;
margin-bottom: 15px;
margin-left: auto;
margin-right: auto;
border-bottom-right-radius: .5rem !important;
border-bottom-left-radius: .5rem !important;
}
#gallery>div>.h-full {
min-height: 20rem;
}
.details:hover {
text-decoration: underline;
}
.gr-button {
white-space: nowrap;
}
.gr-button:focus {
border-color: rgb(147 197 253 / var(--tw-border-opacity));
outline: none;
box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
--tw-border-opacity: 1;
--tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
--tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
--tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
--tw-ring-opacity: .5;
}
#advanced-btn {
font-size: .7rem !important;
line-height: 19px;
margin-top: 12px;
margin-bottom: 12px;
padding: 2px 8px;
border-radius: 14px !important;
}
#advanced-options {
display: none;
margin-bottom: 20px;
}
.footer {
margin-bottom: 45px;
margin-top: 35px;
text-align: center;
border-bottom: 1px solid #e5e5e5;
}
.footer>p {
font-size: .8rem;
display: inline-block;
padding: 0 10px;
transform: translateY(10px);
background: white;
}
.dark .footer {
border-color: #303030;
}
.dark .footer>p {
background: #0b0f19;
}
.acknowledgments h4{
margin: 1.25em 0 .25em 0;
font-weight: bold;
font-size: 115%;
}
.animate-spin {
animation: spin 1s linear infinite;
}
@keyframes spin {
from {
transform: rotate(0deg);
}
to {
transform: rotate(360deg);
}
}
#share-btn-container {
display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
margin-top: 10px;
margin-left: auto;
}
#share-btn {
all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;right:0;
}
#share-btn * {
all: unset;
}
#share-btn-container div:nth-child(-n+2){
width: auto !important;
min-height: 0px !important;
}
#share-btn-container .wrap {
display: none !important;
}
.gr-form{
flex: 1 1 50%; border-top-right-radius: 0; border-bottom-right-radius: 0;
}
#prompt-container{
gap: 0;
}
#prompt-text-input, #negative-prompt-text-input{padding: .45rem 0.625rem}
#component-16{border-top-width: 1px!important;margin-top: 1em}
.image_duplication{position: absolute; width: 100px; left: 50px}
.generate-container {display: flex; justify-content: flex-end;}
#generate-btn {background: linear-gradient(to bottom right, #ffedd5, #fdba74)}
"""
block = gr.Blocks(css=css)
# text, negative, width, height, sampler, steps, seed, guidance_scale
# examples = [
# [
# 'A high tech solarpunk utopia in the Amazon rainforest',
# 'low quality',
# 512,
# 512,
# 'ddim',
# 30,
# 0,
# 9
# ],
# [
# 'A pikachu fine dining with a view to the Eiffel Tower',
# 'low quality',
# 512,
# 512,
# 'ddim',
# 30,
# 0,
# 9
# ],
# [
# 'A mecha robot in a favela in expressionist style',
# 'low quality, 3d, photorealistic',
# 512,
# 512,
# 'ddim',
# 30,
# 0,
# 9
# ],
# [
# 'an insect robot preparing a delicious meal',
# 'low quality, illustration',
# 512,
# 512,
# 'ddim',
# 30,
# 0,
# 9
# ],
# [
# "A small cabin on top of a snowy mountain in the style of Disney, artstation",
# 'low quality, ugly',
# 512,
# 512,
# 'ddim',
# 30,
# 0,
# 9
# ],
# ]
examples = list(map(lambda x: [
x,
'low quality',
512,
512,
'DPMSolverMultistep',
30,
0,
9
], word_list))[:500]
with block:
title = "Stable Diffusion 2.1 Demo"
desc = """ small stable diffusion Demo App.
Click Generate image Button to generate image.
Also Change params to have a try
more size may cost more time.
It's just a simplified demo, you can use more advanced features optimize image quality
"""
tutorial_link = "https://docs.cworld.ai/docs/cworld-ai/quick-start-stable-diffusion"
gr.HTML(
f"""