Update app.py

app.py

@@ -3,7 +3,7 @@ import json
 import gradio as gr
 import numpy as np
 from gradio import processing_utils
-
+import requests
 from packaging import version
 from PIL import Image, ImageDraw
 import functools
@@ -171,28 +171,6 @@ def upload_callback(image_input, state, visual_chatgpt=None):
     return state, state, image_input, click_state, image_input, image_input, image_embedding, \
         original_size, input_size
 
-def store_click(image_input, point_prompt, click_mode, state, click_state, evt: gr.SelectData):
-    click_index = evt.index
-    if point_prompt == 'Positive':
-        coordinate = [click_index[0], click_index[1], 1]
-    else:
-        coordinate = [click_index[0], click_index[1], 0]
-    
-    if click_mode == 'Continuous':
-        click_state[0].append(coordinate)
-    elif click_mode == 'Single':
-        click_state[0] = [coordinate]  # Overwrite with latest click
-
-    return state, click_state
-
-def generate_caption(image_input, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt):
-    last_click = click_state[0][-1]
-    point_prompt = 'Positive' if last_click[2] == 1 else 'Negative'
-    
-    return inference_click(image_input, point_prompt, click_mode, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt, gr.SelectData(index=(last_click[0], last_click[1])))
-    
-
-
 
 
 def inference_click(image_input, point_prompt, click_mode, enable_wiki, language, sentiment, factuality,
@@ -230,14 +208,13 @@ def inference_click(image_input, point_prompt, click_mode, enable_wiki, language
     out = model.inference(image_input, prompt, controls, disable_gpt=True, enable_wiki=enable_wiki, verbose=True, args={'clip_filter': False})[0]
 
     state = state + [("Image point: {}, Input label: {}".format(prompt["input_point"], prompt["input_label"]), None)]
-    state = state + [(None, "raw_caption: {}".format(out['generated_captions']['raw_caption']))]
     update_click_state(click_state, out['generated_captions']['raw_caption'], click_mode)
     text = out['generated_captions']['raw_caption']
     input_mask = np.array(out['mask'].convert('P'))
     image_input = mask_painter(np.array(image_input), input_mask)
     origin_image_input = image_input
-    image_input = create_bubble_frame(image_input, text, (click_index[0], click_index[1]), input_mask,
-                                      input_points=input_points, input_labels=input_labels)
+    # image_input = create_bubble_frame(image_input, None, (click_index[0], click_index[1]), input_mask,
+    #                                   input_points=input_points, input_labels=input_labels)
     x, y = input_points[-1]
     
     if visual_chatgpt is not None:
@@ -247,19 +224,62 @@ def inference_click(image_input, point_prompt, click_mode, enable_wiki, language
         point_prompt = f'You should primarly use tools on the selected regional image (description: {text}, path: {new_crop_save_path}), which is a part of the whole image (path: {visual_chatgpt.current_image}). If human mentioned some objects not in the selected region, you can use tools on the whole image.'
         visual_chatgpt.point_prompt = point_prompt
 
-    yield state, state, click_state, image_input
+
+    generated_caption = text
+    print(generated_caption)
+
+    yield state, state, click_state, image_input, generated_caption
+
     if not args.disable_gpt and model.text_refiner:
         refined_caption = model.text_refiner.inference(query=text, controls=controls, context=out['context_captions'],
                                                        enable_wiki=enable_wiki)
-        # new_cap = 'Original: ' + text + '. Refined: ' + refined_caption['caption']
         new_cap = refined_caption['caption']
         if refined_caption['wiki']:
             state = state + [(None, "Wiki: {}".format(refined_caption['wiki']))]
         state = state + [(None, f"caption: {new_cap}")]
-        refined_image_input = create_bubble_frame(origin_image_input, new_cap, (click_index[0], click_index[1]),
-                                                  input_mask,
-                                                  input_points=input_points, input_labels=input_labels)
-        yield state, state, click_state, refined_image_input
+        # refined_image_input = create_bubble_frame(origin_image_input, None, (click_index[0], click_index[1]),
+        #                                           input_mask,
+        #                                           input_points=input_points, input_labels=input_labels)
+        yield state, state, click_state, image_input, new_cap
+
+def submit_caption(image_input, state,generated_caption):
+    print(state)
+    if state and isinstance(state[-1][1], dict):
+        params = state[-1][1]
+    else:
+        params = {}
+
+    click_index = params.get("click_index", (0, 0))
+    input_mask = params.get("input_mask", np.zeros((1, 1)))
+    input_points = params.get("input_points", [])
+    input_labels = params.get("input_labels", [])
+
+    click_index = params.get("click_index", (0, 0))
+    input_mask = params.get("input_mask", np.zeros((1, 1)))
+    input_points = params.get("input_points", [])
+    input_labels = params.get("input_labels", [])
+
+    image_input = create_bubble_frame(np.array(image_input), generated_caption, (click_index[0], click_index[1]), input_mask,
+                                      input_points=input_points, input_labels=input_labels)
+
+    
+    if generated_caption:
+        state = state + [(None, f"RAW_Caption: {generated_caption}")]
+        txt2speech(generated_caption)
+
+    yield state,state,image_input
+
+
+def txt2speech(text): 
+    print("Initializing text-to-speech conversion...") 
+    # API_URL = "https://api-inference.huggingface.co/models/espnet/kan-bayashi_ljspeech_vits" 
+    # headers = {"Authorization": f"Bearer {os.environ['HUGGINGFACEHUB_API_TOKEN']}"} 
+    # payloads = {'inputs': text} 
+    # response = requests.post(API_URL, headers=headers, json=payloads) 
+    # with open('audio_story.mp3', 'wb') as file: 
+    #     file.write(response.content)
+    print("Text-to-speech conversion completed.")
+
 
 
 def get_sketch_prompt(mask: Image.Image):
@@ -322,7 +342,7 @@ def inference_traject(sketcher_image, enable_wiki, language, sentiment, factuali
     origin_image_input = image_input
 
     fake_click_index = (int((boxes[0][0] + boxes[0][2]) / 2), int((boxes[0][1] + boxes[0][3]) / 2))
-    image_input = create_bubble_frame(image_input, text, fake_click_index, input_mask)
+    image_input = create_bubble_frame(image_input, "", fake_click_index, input_mask)
 
     yield state, state, image_input
 
@@ -415,6 +435,7 @@ def create_ui():
         visual_chatgpt = gr.State(None)
         original_size = gr.State(None)
         input_size = gr.State(None)
+        generated_caption = gr.State("")
         # img_caption = gr.State(None)
         aux_state = gr.State([])
 
@@ -442,7 +463,7 @@ def create_ui():
                             with gr.Row(scale=0.4):
                                 clear_button_click = gr.Button(value="Clear Clicks", interactive=True)
                                 clear_button_image = gr.Button(value="Clear Image", interactive=True)
-                                submit_button_click = gr.Button(value="Submit", interactive=True)
+                                submit_button_click=gr.Button(value="Submit", interactive=True)
                     with gr.Tab("Trajectory (beta)"):
                         sketcher_input = ImageSketcher(type="pil", interactive=True, brush_radius=20,
                                                        elem_id="image_sketcher")
@@ -566,8 +587,6 @@ def create_ui():
         )
         clear_button_text.click(clear_chat_memory, inputs=[visual_chatgpt])
         
-        submit_button_click.click(generate_caption, inputs=[origin_image, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt], outputs=[chatbot, state, click_state, image_input])
-
         image_input.clear(
             lambda: (None, [], [], [[], [], []], "", "", ""),
             [],
@@ -595,8 +614,25 @@ def create_ui():
                              [chatbot, state, origin_image, click_state, image_input, sketcher_input,
                               image_embedding, original_size, input_size])
         example_image.change(clear_chat_memory, inputs=[visual_chatgpt])
-        # select coordinate
-        image_input.select(store_click, inputs=[origin_image, point_prompt, click_mode, state, click_state], outputs=[state, click_state])
+
+        image_input.select(
+            inference_click,
+            inputs=[
+                origin_image, point_prompt, click_mode, enable_wiki, language, sentiment, factuality, length,
+                image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt
+            ],
+            outputs=[chatbot, state, click_state, image_input, generated_caption],
+            show_progress=False, queue=True
+        )
+
+        submit_button_click.click(
+            submit_caption,
+            inputs=[image_input, state, generated_caption],
+            outputs=[chatbot,state,image_input],
+            show_progress=True, queue=True
+        )
+
+
 
         submit_button_sketcher.click(
             inference_traject,