svjack's picture
Update simple_app.py
d43d940
import gradio as gr
from Lex import *
'''
lex = Lexica(query="man woman fire snow").images()
'''
from PIL import Image
import imagehash
import requests
from time import sleep
sleep_time = 0.5
hash_func_name = list(filter(lambda x: x.endswith("hash") and
"hex" not in x ,dir(imagehash)))
hash_func_name = ['average_hash', 'colorhash', 'dhash', 'phash', 'whash', 'crop_resistant_hash',]
def min_dim_to_size(img, size = 512):
h, w = img.size
ratio = size / max(h, w)
h, w = map(lambda x: int(x * ratio), [h, w])
return ( ratio ,img.resize((h, w)) )
#ratio_size = 512
#ratio, img_rs = min_dim_to_size(img, ratio_size)
def image_click(images, evt: gr.SelectData):
img_selected = images[evt.index]
return images[evt.index]['name']
def swap_gallery(im, images, func_name):
#### name data is_file
#print(images[0].keys())
if im is None:
return list(map(lambda x: x["name"], images))
hash_func = getattr(imagehash, func_name)
im_hash = hash_func(Image.fromarray(im))
t2_list = sorted(images, key = lambda imm:
hash_func(Image.open(imm["name"])) - im_hash, reverse = False)
return list(map(lambda x: x["name"], t2_list))
def lexica(prompt, limit_size = 128, ratio_size = 256 + 128):
lex = Lexica(query=prompt).images()
lex = lex[:limit_size]
lex = list(map(lambda x: x.replace("full_jpg", "sm2"), lex))
lex_ = []
for ele in lex:
try:
im = Image.open(
requests.get(ele, stream = True).raw
)
lex_.append(im)
except:
print("err")
sleep(sleep_time)
assert lex_
lex = list(map(lambda x: min_dim_to_size(x, ratio_size)[1], lex_))
return lex
def enterpix(prompt, limit_size = 100, ratio_size = 256 + 128, use_key = "bigThumbnailUrl"):
resp = requests.post(
url = "https://www.enterpix.app/enterpix/v1/image/prompt-search",
data= {
"length": limit_size,
"platform": "stable-diffusion,midjourney",
"prompt": prompt,
"start": 0
}
)
resp = resp.json()
resp = list(map(lambda x: x[use_key], resp["images"]))
lex_ = []
for ele in resp:
try:
im = Image.open(
requests.get(ele, stream = True).raw
)
lex_.append(im)
except:
print("err")
sleep(sleep_time)
assert lex_
resp = list(map(lambda x: min_dim_to_size(x, ratio_size)[1], lex_))
return resp
def search(prompt, search_name):
if search_name == "lexica":
return lexica(prompt)
else:
return enterpix(prompt)
with gr.Blocks(css="custom.css") as demo:
title = gr.HTML(
"""<h1><img src="https://i.imgur.com/52VJ8vS.png" alt="SD"> StableDiffusion Image Search </h1>""",
elem_id="title",
)
with gr.Row():
with gr.Column():
with gr.Row():
search_func_name = gr.Radio(choices=["lexica", "enterpix"],
value="lexica", label="Search by", elem_id="search_radio")
with gr.Row():
#inputs = gr.Textbox(label = 'Enter prompt to search Lexica.art')
inputs = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=20, min_width = 256,
placeholder="Enter prompt to search", elem_id="prompt")
#gr.Slider(label='Number of images ', minimum = 4, maximum = 20, step = 1, value = 4)]
text_button = gr.Button("Retrieve Images", elem_id="run_button")
i = gr.Image(elem_id="result-image")
gr.Examples(
[
["Chinese Traditional Culture", "lexica", "images/AbsoluteReality.png"],
["trending digital art", "lexica", "images/waifu.png"],
["beautiful home", "enterpix", "images/Cyberpunk_Anime.png"],
["interior design of living room", "enterpix", "images/DreamShaper.png"],
],
inputs = [inputs, search_func_name, i],
label = "Examples"
)
with gr.Column():
title = gr.Markdown(
value="### Click on a Image in the gallery to select it, and the grid order will change",
visible=True,
elem_id="selected_model",
)
order_func_name = gr.Radio(choices=hash_func_name,
value=hash_func_name[0], label="Order by", elem_id="order_radio")
outputs = gr.Gallery(lable='Output gallery', elem_id="gallery",).style(grid=5,height=768 - 128,
allow_preview=False, label = "retrieve Images")
#gr.Dataframe(label='prompts for corresponding images')]
#outputs.select(image_click, outputs, i, _js="(x) => x.splice(0,x.length)")
outputs.select(image_click, outputs, i,)
i.change(
fn=swap_gallery,
inputs=[i, outputs, order_func_name],
outputs=outputs,
queue=False
)
order_func_name.change(
fn=swap_gallery,
inputs=[i, outputs, order_func_name],
outputs=outputs,
queue=False
)
#### gr.Textbox().submit().success()
### lexica
#text_button.click(lexica, inputs=inputs, outputs=outputs)
### enterpix
#text_button.click(enterpix, inputs=inputs, outputs=outputs)
text_button.click(search, inputs=[inputs, search_func_name], outputs=outputs)
demo.launch("0.0.0.0")