added diffusion functionality

app.py

@@ -46,6 +46,8 @@ pipe = StableDiffusionPipeline.from_pretrained(
 )
 pipe.to(device)
 
+missing_im = Image.open('missing.png')
+
 def encode_text(text: str):
     text_inputs = pipe.tokenizer(
         text, return_tensors='pt'
@@ -64,28 +66,68 @@ def prompt_query(text: str):
     prompts = list(dict.fromkeys(prompts))
     return [[x] for x in prompts[:5]]
 
+def diffuse(text: str):
+    # diffuse
+    out = pipe(text)
+    if any(out.nsfw_content_detected):
+        return {}
+    else:
+        _id = str(uuid.uuid4())
+        # add image to Cloud Storage
+        im = out.images[0]
+        im.save(f'{_id}.png', format='png')
+        # push to storage
+        blob = bucket.blob(f'images/{_id}.png')
+        blob.upload_from_filename(f'{_id}.png')
+        # delete local file
+        os.remove(f'{_id}.png')
+        # add embedding and metadata to Pinecone
+        embeds = encode_text(text)
+        meta = {
+            'prompt': text,
+            'image_url': f'images/{_id}.png'
+        }
+        index.upsert([(_id, embeds, meta)])
+        return out.images[0]
+
 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()
+        print(f"DELETED '{_id}'")
+        return False
+
 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']
     ]
+    scores = [match['score'] for match in xc['matches']]
+    ids = [match['id'] for match in xc['matches']]
     images = []
-    for image_url in image_urls:
+    for _id, image_url in zip(ids, image_urls):
         try:
             blob = bucket.blob(image_url).download_as_string()
             blob_bytes = io.BytesIO(blob)
             im = Image.open(blob_bytes)
-            images.append(im)
+            if test_image(_id, im):
+                images.append(im)
+            else:
+                images.append(missing_im)
         except ValueError:
-            print(f"error for '{image_url}'")
-    return images
+            print(f"ValueError: '{image_url}'")
+    return images, scores
 
 # __APP FUNCTIONS__
 
@@ -93,8 +135,19 @@ 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)
+    images, scores = prompt_image(text)
+    match_found = False
+    for score in scores:
+        if score > 0.85:
+            match_found = True
+    if match_found:
+        print("MATCH FOUND")
+        return gr.Gallery.update(value=images)
+    else:
+        print("NO MATCH FOUND")
+        diffuse(text)
+        images, scores = prompt_image(text)
+        return gr.Gallery.update(value=images)
 
 # __CREATE APP__
 demo = gr.Blocks()
@@ -129,12 +182,14 @@ with demo:
                 suggestions.components
             )
             
-        
         # results column
         with gr.Column():
             pics = gr.Gallery()
             pics.style(grid=3)
             # search event listening
-            search.click(set_images, prompt, pics)
+            try:
+                search.click(set_images, prompt, pics)
+            except OSError:
+                print("OSError")
 
 demo.launch()