Update app.py

app.py

@@ -1,4 +1,5 @@
 from io import BytesIO
+import io
 from math import inf
 import os
 import base64
@@ -26,7 +27,15 @@ import easyocr
 import re
 import edge_tts
 from langchain import __version__
-
+import torch
+import gradio as gr
+from transformers import AutoProcessor, SiglipModel
+import faiss
+from huggingface_hub import hf_hub_download
+from datasets import load_dataset
+import pandas as pd
+import requests
+import spaces
 # Print the current version of LangChain
 print(f"Current LangChain version: {__version__}")
 # import tts 
@@ -37,7 +46,9 @@ print(f"Current LangChain version: {__version__}")
 
 
 # import spaces  #
+import threading
 
+lock = threading.Lock()
 import os
 # import uuid
 # from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
@@ -304,6 +315,56 @@ def make3d(images):
 ###############################################################################
 
 
+###############################################################################
+############# This part is for sCLIP #############
+###############################################################################
+
+# download model and dataset
+hf_hub_download("merve/siglip-faiss-wikiart", "siglip_10k_latest.index", local_dir="./")
+hf_hub_download("merve/siglip-faiss-wikiart", "wikiart_10k_latest.csv", local_dir="./")
+
+# read index, dataset and load siglip model and processor
+index = faiss.read_index("./siglip_10k_latest.index")
+df = pd.read_csv("./wikiart_10k_latest.csv")
+device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
+processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")
+slipmodel = SiglipModel.from_pretrained("google/siglip-base-patch16-224").to(device)
+
+
+def read_image_from_url(url):
+    response = requests.get(url)
+    img = Image.open(BytesIO(response.content)).convert("RGB")
+    return img
+
+#@spaces.GPU
+def extract_features_siglip(image):
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors="pt").to(device)
+        image_features = slipmodel.get_image_features(**inputs)
+    return image_features
+
+@spaces.GPU
+def infer(image_path):
+  input_image = Image.open(image_path).convert("RGB")
+  input_features = extract_features_siglip(input_image.convert("RGB"))
+  input_features = input_features.detach().cpu().numpy()
+  input_features = np.float32(input_features)
+  faiss.normalize_L2(input_features)
+  distances, indices = index.search(input_features, 3)
+  gallery_output = []
+  for i,v in enumerate(indices[0]):
+    sim = -distances[0][i]
+    image_url = df.iloc[v]["Link"]
+    img_retrieved = read_image_from_url(image_url)
+    gallery_output.append(img_retrieved)
+  return gallery_output
+
+
+###############################################################################
+############# Above part is for sCLIP #############
+###############################################################################
+
+
 ###############################################################################
 ############# this part is for text to image #############
 ###############################################################################
@@ -623,6 +684,14 @@ async def chat_input_callback(*args):
 
         
 def upload_callback(image_input, state, visual_chatgpt=None, openai_api_key=None,language="English"):
+    if isinstance(image_input, dict):  # if upload from sketcher_input, input contains image and mask
+        image_input = image_input['background']
+        
+    if isinstance(image_input, str):
+        image_input = Image.open(io.BytesIO(base64.b64decode(image_input)))
+    elif isinstance(image_input, bytes):
+        image_input = Image.open(io.BytesIO(image_input))
+
     
     click_state = [[], [], []]
     image_input = image_resize(image_input, res=1024)
@@ -938,13 +1007,15 @@ async def inference_traject(origin_image,sketcher_image, enable_wiki, language,
     
     if trace_type=="Trace+Seg":
         input_mask = np.array(out['mask'].convert('P'))
-        image_input = mask_painter(np.array(image_input), input_mask, background_alpha=0 )
+        image_input = mask_painter(np.array(image_input), input_mask, background_alpha=0)
+        d3_input=mask_painter(np.array(image_input), input_mask)
         crop_save_path=out['crop_save_path']
     
     else:
         image_input = Image.fromarray(np.array(origin_image))
         draw = ImageDraw.Draw(image_input)
         draw.rectangle(boxes, outline='red', width=2) 
+        d3_input=image_input
         cropped_image = origin_image.crop(boxes)
         cropped_image.save('temp.png')
         crop_save_path='temp.png'
@@ -977,14 +1048,14 @@ async def inference_traject(origin_image,sketcher_image, enable_wiki, language,
         try:
             audio_output = await texttospeech(read_info, language,autoplay)
             # return state, state, refined_image_input, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state, waveform_visual, audio_output
-            return state, state,image_input,audio_output
+            return state, state,image_input,audio_output,crop_save_path,d3_input
 
 
         except Exception as e:
             state = state + [(None, f"Error during TTS prediction: {str(e)}")]
             print(f"Error during TTS prediction: {str(e)}")
             # return state, state, refined_image_input, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state, None, None
-            return state, state, image_input,audio_output
+            return state, state, image_input,audio_output,crop_save_path,d3_input
 
 
     else:
@@ -1290,11 +1361,10 @@ def create_ui():
 
         with gr.Row():                        
             
-            with gr.Column():
+            with gr.Column(scale=6):
                 with gr.Column(visible=False) as modules_not_need_gpt:
                     with gr.Tab("Base(GPT Power)") as base_tab:
                         image_input_base = gr.Image(type="pil", interactive=True, elem_id="image_upload")
-                        example_image = gr.Image(type="pil", interactive=False, visible=False)
                         with gr.Row():
                             name_label_base = gr.Button(value="Name: ")
                             artist_label_base = gr.Button(value="Artist: ")
@@ -1304,45 +1374,51 @@ def create_ui():
                     with gr.Tab("Click") as click_tab:
                         image_input = gr.Image(type="pil", interactive=True, elem_id="image_upload")
                         example_image = gr.Image(type="pil", interactive=False, visible=False)
+                        # example_image_click = gr.Image(type="pil", interactive=False, visible=False)
                         with gr.Row():
                             name_label = gr.Button(value="Name: ")
                             artist_label = gr.Button(value="Artist: ")
                             year_label = gr.Button(value="Year: ")
                             material_label = gr.Button(value="Material: ")
                         with gr.Row():
-                            with gr.Row():
-                                focus_type = gr.Radio(
-                                        choices=["CFV-D", "CFV-DA", "CFV-DAI","PFV-DDA"],
-                                        value="CFV-D",
-                                        label="Information Type",
+                            with gr.Column():
+                                with gr.Row():
+                                    focus_type = gr.Radio(
+                                            choices=["CFV-D", "CFV-DA", "CFV-DAI","PFV-DDA"],
+                                            value="CFV-D",
+                                            label="Information Type",
+                                            interactive=True,
+                                            scale=4)
+                                
+                                with gr.Row():
+                                    point_prompt = gr.Radio(
+                                        choices=["Positive", "Negative"],
+                                        value="Positive",
+                                        label="Point Prompt",
+                                        scale=5,
                                         interactive=True)
-                            with gr.Row():
-                                submit_button_click=gr.Button(value="Submit", interactive=True,variant='primary',size="sm")
-                        with gr.Row():
-                            with gr.Row():
-                                point_prompt = gr.Radio(
-                                    choices=["Positive", "Negative"],
-                                    value="Positive",
-                                    label="Point Prompt",
-                                    interactive=True)
-                                click_mode = gr.Radio(
-                                    choices=["Continuous", "Single"],
-                                    value="Continuous",
-                                    label="Clicking Mode",
-                                    interactive=True)
-                            with gr.Row():
-                                clear_button_click = gr.Button(value="Clear Clicks", interactive=True)
-                                clear_button_image = gr.Button(value="Clear Image", interactive=True)
+                                    click_mode = gr.Radio(
+                                        choices=["Continuous", "Single"],
+                                        value="Continuous",
+                                        label="Clicking Mode",
+                                        scale=5,
+                                        interactive=True)
+                            with gr.Column():
+                                with gr.Row():
+                                    submit_button_click=gr.Button(value="Submit", interactive=True,variant='primary',scale=2)
+                                with gr.Row():
+                                    clear_button_click = gr.Button(value="Clear Clicks", interactive=True,scale=2)
+                                    clear_button_image = gr.Button(value="Clear Image", interactive=True,scale=2)
                                 
                     with gr.Tab("Trajectory (beta)") as traj_tab:
                         # sketcher_input = ImageSketcher(type="pil", interactive=True, brush_radius=10, 
                         #                                elem_id="image_sketcher")
                         sketcher_input = gr.ImageEditor(type="pil", interactive=True, 
                                                        elem_id="image_sketcher")
-                        example_image = gr.Image(type="pil", interactive=False, visible=False)
+                        # example_image_traj = gr.Image(type="pil", interactive=False, visible=False)
                         with gr.Row():                   
-                            submit_button_sketcher = gr.Button(value="Submit", interactive=True)
                             clear_button_sketcher = gr.Button(value="Clear Sketch", interactive=True)
+                            submit_button_sketcher = gr.Button(value="Submit", interactive=True)
                         with gr.Row():
                             with gr.Row():
                                 focus_type_sketch = gr.Radio(
@@ -1354,9 +1430,9 @@ def create_ui():
                                         choices=["Trace+Seg", "Trace"],
                                         value="Trace+Seg",
                                         label="Trace Type",
-                                        interactive=True)
+                                        interactive=True)                    
 
-                    with gr.Column(visible=False) as modules_need_gpt1:
+                    with gr.Column(visible=False,scale=4) as modules_need_gpt1:
                         with gr.Row():
                             sentiment = gr.Radio(
                                 choices=["Positive", "Natural", "Negative"],
@@ -1395,7 +1471,7 @@ def create_ui():
                 
                 
 
-            with gr.Column():                    
+            with gr.Column(scale=5):                    
                 with gr.Column(visible=True) as module_key_input:
                     openai_api_key = gr.Textbox(
                         placeholder="Input openAI API key",
@@ -1454,7 +1530,7 @@ def create_ui():
 
             with gr.Column():
                 with gr.Column():
-                    gr.Radio([artist], label="Artist", info="Who is the artist?🧑‍🎨"),
+                    gr.Radio(["Other Paintings by the Artist"], label="Artist", info="Who is the artist?🧑‍🎨"),
                     gr.Radio(["Oil Painting","Printmaking","Watercolor Painting","Drawing"], label="Art Forms", info="What are the art forms?🎨"),
                     gr.Radio(["Renaissance", "Baroque", "Impressionism","Modernism"], label="Period", info="Which art period?⏳"),
                     # to be done 
@@ -1582,20 +1658,9 @@ def create_ui():
         #     api_name="run",
         # )
         run_button.click(
-            fn=generate,
-            inputs=[
-                prompt,
-                negative_prompt,
-                use_negative_prompt,
-                seed,
-                width,
-                height,
-                guidance_scale,
-                num_inference_steps,
-                randomize_seed,
-                num_images
-            ],
-            outputs=[result, seed]
+            fn=infer,
+            inputs=[new_crop_save_path],
+            outputs=[result]
             )
         
         ###############################################################################
@@ -1825,12 +1890,12 @@ def create_ui():
                            [chatbot, state, origin_image, click_state, image_input, image_input_base, sketcher_input,
                             image_embedding, original_size, input_size,name_label,artist_label,year_label,material_label,name_label_base, artist_label_base, year_label_base, material_label_base,paragraph,artist])
         
-        image_input.upload(upload_callback, [image_input, state, visual_chatgpt, openai_api_key],
-                           [chatbot, state, origin_image, click_state, image_input, image_input_base, sketcher_input,
-                            image_embedding, original_size, input_size,name_label,artist_label,year_label,material_label,name_label_base, artist_label_base, year_label_base, material_label_base,paragraph,artist])
-        sketcher_input.upload(upload_callback, [sketcher_input, state, visual_chatgpt, openai_api_key],
-                              [chatbot, state, origin_image, click_state, image_input, image_input_base, sketcher_input,
-                               image_embedding, original_size, input_size,name_label,artist_label,year_label,material_label,name_label_base, artist_label_base, year_label_base, material_label_base,paragraph,artist])
+        # image_input.upload(upload_callback, [image_input, state, visual_chatgpt, openai_api_key],
+        #                    [chatbot, state, origin_image, click_state, image_input, image_input_base, sketcher_input,
+        #                     image_embedding, original_size, input_size,name_label,artist_label,year_label,material_label,name_label_base, artist_label_base, year_label_base, material_label_base,paragraph,artist])
+        # sketcher_input.upload(upload_callback, [sketcher_input, state, visual_chatgpt, openai_api_key],
+        #                       [chatbot, state, origin_image, click_state, image_input, image_input_base, sketcher_input,
+        #                        image_embedding, original_size, input_size,name_label,artist_label,year_label,material_label,name_label_base, artist_label_base, year_label_base, material_label_base,paragraph,artist])
         chat_input.submit(chat_input_callback, [visual_chatgpt, chat_input, click_state, state, aux_state,language,auto_play],
                           [chatbot, state, aux_state,output_audio])
         chat_input.submit(lambda: "", None, chat_input)
@@ -1904,7 +1969,7 @@ def create_ui():
                 origin_image,sketcher_input, enable_wiki, language, sentiment, factuality, length, image_embedding, state,
                 original_size, input_size, text_refiner,focus_type_sketch,paragraph,openai_api_key,auto_play,Input_sketch
             ],
-            outputs=[chatbot, state, sketcher_input,output_audio],
+            outputs=[chatbot, state, sketcher_input,output_audio,new_crop_save_path,input_image],
             show_progress=False, queue=True
         )