Create app.py

app.py

@@ -0,0 +1,144 @@
+from diffusers import StableDiffusionPipeline
+import torch
+import io
+from PIL import Image
+import os
+from google.cloud import storage
+import pinecone
+
+# create Storage Cloud credentials
+G_API = {
+    "type": os.environ["type"],
+    "project_id": os.environ["project_id"],
+    "private_key_id": os.environ["private_key_id"],
+    "private_key": os.environ["private_key"],
+    "client_email": os.environ["client_email"],
+    "client_id": os.environ["client_id"],
+    "auth_uri": os.environ["auth_uri"],
+    "token_uri": os.environ["token_uri"],
+    "auth_provider_x509_cert_url": os.environ["auth_provider_x509_cert_url"],
+    "client_x509_cert_url": os.environ["client_x509_cert_url"]
+}
+with open('cloud-storage.json', 'w', encoding='utf-8') as fp:
+    fp.write(json.dumps(G_API))
+del G_API
+# connect to Cloud Storage
+os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'cloud-storage.json'
+storage_client = storage.Client()
+bucket = storage_client.get_bucket('hf-diffusion-images')
+    
+# get api key for pinecone auth
+PINECONE_KEY = os.environ['PINECONE_KEY']
+
+index_id = "hf-diffusion"
+
+# 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 = 'cpu'
+
+# init all of the models and move them to a given GPU
+pipe = StableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4", use_auth_token=True
+)
+pipe.to(device)
+
+def encode_text(text: str):
+    text_inputs = pipe.tokenizer(
+        text, return_tensors='pt'
+    ).to(device)
+    text_embeds = pipe.text_encoder(**text_inputs)
+    text_embeds = text_embeds.pooler_output.cpu().tolist()[0]
+    return text_embeds
+
+def prompt_query(text: str):
+    embeds = encode_text(text)
+    xc = index.query(embeds, top_k=30, include_metadata=True)
+    prompts = [
+        match['metadata']['prompt'] for match in xc['matches']
+    ]
+    # deduplicate while preserving order
+    prompts = list(dict.fromkeys(prompts))
+    return [[x] for x in prompts[:5]]
+
+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 prompt_image(text: str):
+    embeds = encode_text(text)
+    xc = index.query(embeds, top_k=9, include_metadata=True)
+    image_urls = [
+        match['metadata']['image_url'] for match in xc['matches']
+    ]
+    images = []
+    for image_url in image_urls:
+        try:
+            blob = bucket.blob(image_url).download_as_string()
+            blob_bytes = io.BytesIO(blob)
+            im = Image.open(blob_bytes)
+            images.append(im)
+        except ValueError:
+            print(f"error for '{image_url}'")
+    return images
+
+# __APP FUNCTIONS__
+
+def set_suggestion(text: str):
+    return gr.TextArea.update(value=text[0])
+
+def set_images(text: str):
+    images = prompt_image(text)
+    return gr.Gallery.update(value=images)
+
+# __CREATE APP__
+demo = gr.Blocks()
+
+with demo:
+    gr.Markdown(
+        """
+        # Dream Cacher
+        """
+    )
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.TextArea(
+                value="A dream about a cat",
+                placeholder="Enter a prompt to dream about",
+                interactive=True
+            )
+            search = gr.Button(value="Search!")
+            suggestions = gr.Dataset(
+                components=[prompt],
+                samples=[
+                    ["Something"],
+                    ["something else"]
+                ]
+            )
+            # event listener for change in prompt
+            prompt.change(prompt_query, prompt, suggestions)
+            # event listener for click on suggestion
+            suggestions.click(
+                set_suggestion,
+                suggestions,
+                suggestions.components
+            )
+            
+        
+        # results column
+        with gr.Column():
+            pics = gr.Gallery()
+            pics.style(grid=3)
+            # search event listening
+            search.click(set_images, prompt, pics)
+
+demo.launch()