Spaces:
Runtime error
Runtime error
File size: 6,393 Bytes
2118bfe d3ba2a4 2118bfe d3ba2a4 2118bfe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
import gradio as gr
from transformers import CLIPTokenizerFast, CLIPProcessor, CLIPModel
import torch
import io
from PIL import Image
import os
from cryptography.fernet import Fernet
from google.cloud import storage
import pinecone
import json
# decrypt Storage Cloud credentials
fernet = Fernet(os.environ['DECRYPTION_KEY'])
with open('cloud-storage.encrypted', 'rb') as fp:
encrypted = fp.read()
creds = json.loads(fernet.decrypt(encrypted).decode())
# then save creds to file
with open('cloud-storage.json', 'w', encoding='utf-8') as fp:
fp.write(json.dumps(creds, indent=4))
# connect to Cloud Storage
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'cloud-storage.json'
storage_client = storage.Client()
bucket = storage_client.get_bucket('diffusion-search')
# get api key for pinecone auth
PINECONE_KEY = os.environ['PINECONE_KEY']
index_id = "diffusion-search"
# init connection to pinecone
pinecone.init(
api_key=PINECONE_KEY,
environment="us-west1-gcp"
)
if index_id not in pinecone.list_indexes():
raise ValueError(f"Index '{index_id}' not found")
index = pinecone.Index(index_id)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using '{device}' device...")
# init all of the models and move them to a given GPU
# if you have CUDA or MPS, set it to the active device like this
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "openai/clip-vit-base-patch32"
# we initialize a tokenizer, image processor, and the model itself
tokenizer = CLIPTokenizerFast.from_pretrained(model_id)
model = CLIPModel.from_pretrained(model_id).to(device)
missing_im = Image.open('missing.png')
threshold = 0.85
def encode_text(text: str):
# create transformer-readable tokens
inputs = tokenizer(text, return_tensors="pt").to(device)
text_emb = model.get_text_features(**inputs).cpu().detach().tolist()
return text_emb
def prompt_query(text: str):
print(f"Running prompt_query('{text}')")
embeds = encode_text(text)
try:
xc = index.query(embeds, top_k=30, include_metadata=True)
except Exception as e:
print(f"Error during query: {e}")
# reinitialize connection
pinecone.init(api_key=PINECONE_KEY, environment='us-west1-gcp')
index2 = pinecone.Index(index_id)
try:
xc = index2.query(embeds, top_k=30, include_metadata=True)
print("Reinitialized query successful")
except Exception as e:
raise ValueError(e)
scores = [round(match['score'], 2) for match in xc['matches']]
ids = [match['id'] for match in xc['matches']]
return ids
def get_image(url: str):
blob = bucket.blob(url).download_as_string()
blob_bytes = io.BytesIO(blob)
im = Image.open(blob_bytes)
return im
def test_image(_id, image):
try:
image.save('tmp.png')
return True
except OSError:
# delete corrupted file from pinecone and cloud
index.delete(ids=[_id])
bucket.blob(f"images/{_id}.png").delete()
return False
def prompt_image(text: str):
embeds = encode_text(text)
try:
xc = index.query(
embeds, top_k=9, include_metadata=True,
filter={"image_nsfw": {"$lt": 0.5}}
)
except Exception as e:
print(f"Error during query: {e}")
# reinitialize connection
pinecone.init(api_key=PINECONE_KEY, environment='us-west1-gcp')
index2 = pinecone.Index(index_id)
try:
xc = index2.query(
embeds, top_k=9, include_metadata=True,
filter={"image_nsfw": {"$lt": 0.5}}
)
print("Reinitialized query successful")
except Exception as e:
raise ValueError(e)
scores = [match['score'] for match in xc['matches']]
ids = [match['id'] for match in xc['matches']]
images = []
for _id in ids:
try:
image_url = f"images/{_id}.png"
blob = bucket.blob(image_url).download_as_string()
blob_bytes = io.BytesIO(blob)
im = Image.open(blob_bytes)
if test_image(_id, im):
images.append(im)
else:
images.append(missing_im)
except ValueError:
print(f"ValueError: '{image_url}'")
return images, scores
# __APP FUNCTIONS__
def set_suggestion(text: str):
return gr.TextArea.update(value=text[0])
def set_images(text: str):
images, scores = prompt_image(text)
return gr.Gallery.update(value=images)
# __CREATE APP__
demo = gr.Blocks()
with demo:
gr.HTML(
"""
<img src="https://huggingface.co/spaces/pinecone/diffusion-image-search/resolve/main/pine-trees-collage.png" />
<style>
.parallax {
/* The image used */
background-image: url("https://huggingface.co/spaces/pinecone/diffusion-image-search/resolve/main/pine-trees-collage.png");
/* Create the parallax scrolling effect */
background-attachment: fixed;
background-position: center;
background-repeat: no-repeat;
background-size: cover;
}
</style>
<!-- Container element -->
<div class="parallax"></div>
"""
)
with gr.Row():
with gr.Column():
prompt = gr.TextArea(
value="space dogs",
placeholder="Something cool to search for...",
interactive=True
)
search = gr.Button(value="Search!")
gr.Markdown(
"""
#### Search through 10K images generated by AI
This app demonstrates the idea of text-to-image search. The search process
uses an AI model that understands the *meaning* of text and images to identify
images that best align to a search prompt.
πͺ [*Built with the OP Stack*](https://gkogan.notion.site/gkogan/The-OP-Stack-aafcab0005e3445a8ad8491aac80446c)
"""
)
# results column
with gr.Column():
pics = gr.Gallery()
pics.style(grid=3)
# search event listening
try:
search.click(set_images, prompt, pics)
except OSError:
print("OSError")
demo.launch() |