Spaces:
Sleeping
Sleeping
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") | |