fix gpt usage

app.py

@@ -7,7 +7,7 @@ import requests
 from packaging import version
 from PIL import Image, ImageDraw
 import functools
-
+from langchain.llms.openai import OpenAI
 from caption_anything.model import CaptionAnything
 from caption_anything.utils.image_editing_utils import create_bubble_frame
 from caption_anything.utils.utils import mask_painter, seg_model_map, prepare_segmenter, image_resize
@@ -68,17 +68,38 @@ def build_caption_anything_with_models(args, api_key="", captioner=None, sam_mod
     return CaptionAnything(args, api_key, captioner=captioner, segmenter=segmenter, ocr_reader=ocr_reader, text_refiner=text_refiner)
 
 
+def validate_api_key(api_key):
+    api_key = str(api_key).strip()
+    print(api_key)
+    try:
+        test_llm = OpenAI(model_name="gpt-3.5-turbo", temperature=0, openai_api_key=api_key)
+        response = test_llm("Test API call")
+        print(response)
+        return True
+    except Exception as e:
+        print(f"API key validation failed: {e}")
+        return False
+
+
 def init_openai_api_key(api_key=""):
     text_refiner = None
     visual_chatgpt = None
     if api_key and len(api_key) > 30:
-        try:
-            text_refiner = build_text_refiner(args.text_refiner, args.device, args, api_key)
-            assert len(text_refiner.llm('hi')) > 0 # test
-            visual_chatgpt = ConversationBot(shared_chatbot_tools, api_key)
-        except:
-            text_refiner = None
-            visual_chatgpt = None
+        print(api_key)
+        if validate_api_key(api_key):
+            try:
+                text_refiner = build_text_refiner(args.text_refiner, args.device, args, api_key)
+                assert len(text_refiner.llm('hi')) > 0 # test
+                visual_chatgpt = ConversationBot(shared_chatbot_tools, api_key)
+            except Exception as e:
+                print(f"Error initializing TextRefiner or ConversationBot: {e}")
+                text_refiner = None
+                visual_chatgpt = None
+        else:
+            print("Invalid API key.")
+    else:
+        print("API key is too short.")
+    print(text_refiner)
     openai_available = text_refiner is not None
     if openai_available:
         return [gr.update(visible=True)]*6 + [gr.update(visible=False)]*2 + [text_refiner, visual_chatgpt, None]
@@ -175,7 +196,7 @@ def upload_callback(image_input, state, visual_chatgpt=None):
 
 def 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,
-                    evt: gr.SelectData):
+                    out_state, click_index_state, input_mask_state, input_points_state, input_labels_state, evt: gr.SelectData):
     click_index = evt.index
 
     if point_prompt == 'Positive':
@@ -212,11 +233,13 @@ def inference_click(image_input, point_prompt, click_mode, enable_wiki, language
     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, None, (click_index[0], click_index[1]), input_mask,
-    #                                   input_points=input_points, input_labels=input_labels)
-    x, y = input_points[-1]
-    
+
+    click_index_state = click_index
+    input_mask_state = input_mask
+    input_points_state = input_points
+    input_labels_state = input_labels
+    out_state = out  
+
     if visual_chatgpt is not None:
         print('inference_click: add caption to chatGPT memory')
         new_crop_save_path = get_new_image_name('chat_image', func_name='crop')
@@ -224,50 +247,59 @@ 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
 
-
     generated_caption = text
     print(generated_caption)
 
-    yield state, state, click_state, image_input, generated_caption
+    yield state, state, click_state, image_input, generated_caption, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state
 
-    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 = 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, 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,
+def submit_caption(image_input, state, generated_caption, text_refiner, visual_chatgpt, enable_wiki, length, sentiment, factuality, language, 
+                   out_state, click_index_state, input_mask_state, input_points_state, input_labels_state):
+    print("state",state)
+
+    click_index = click_index_state
+    input_mask = input_mask_state
+    input_points = input_points_state
+    input_labels = input_labels_state
+    out = out_state
+    print("click",click_index)
+
+    origin_image_input = image_input
+
+    controls = {
+        'length': length,
+        'sentiment': sentiment,
+        'factuality': factuality,
+        'language': language
+    }
+
+    image_input = create_bubble_frame(np.array(image_input), generated_caption, click_index, input_mask,
                                       input_points=input_points, input_labels=input_labels)
 
-    
     if generated_caption:
         state = state + [(None, f"RAW_Caption: {generated_caption}")]
-        txt2speech(generated_caption)
+        
+
+        if not args.disable_gpt and text_refiner:
+            refined_caption = text_refiner.inference(query=generated_caption, controls=controls, context=out['context_captions'], enable_wiki=enable_wiki)
+            new_cap = refined_caption['caption']
+            if refined_caption.get('wiki'):
+                state = state + [(None, "Wiki: {}".format(refined_caption['wiki']))]
+            state = state + [(None, f"RAW_Caption: {new_cap}")]
+            print("new_cap",new_cap)
+            refined_image_input = create_bubble_frame(np.array(origin_image_input), new_cap, click_index, input_mask,
+                                                      input_points=input_points, input_labels=input_labels)
+            txt2speech(new_cap)
+            yield state, state, refined_image_input, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state
+    
+        else:
+            txt2speech(generated_caption)
+            yield state, state, image_input, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state
+
+
 
-    yield state,state,image_input
 
 
 def txt2speech(text): 
@@ -427,8 +459,8 @@ def create_ui():
             css=get_style()
     ) as iface:
         state = gr.State([])
+        out_state = gr.State(None)
         click_state = gr.State([[], [], []])
-        # chat_state = gr.State([])
         origin_image = gr.State(None)
         image_embedding = gr.State(None)
         text_refiner = gr.State(None)
@@ -436,8 +468,11 @@ def create_ui():
         original_size = gr.State(None)
         input_size = gr.State(None)
         generated_caption = gr.State("")
-        # img_caption = gr.State(None)
         aux_state = gr.State([])
+        click_index_state = gr.State((0, 0))
+        input_mask_state = gr.State(np.zeros((1, 1)))
+        input_points_state = gr.State([])
+        input_labels_state = gr.State([])
 
         gr.Markdown(title)
         gr.Markdown(description)
@@ -619,18 +654,22 @@ def create_ui():
             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
+                image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt,
+                out_state, click_index_state, input_mask_state, input_points_state, input_labels_state
             ],
-            outputs=[chatbot, state, click_state, image_input, generated_caption],
+            outputs=[chatbot, state, click_state, image_input, generated_caption, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state],
             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_caption,
+        inputs=[image_input, state, generated_caption, text_refiner, visual_chatgpt, enable_wiki, length, sentiment, factuality, language, 
+                out_state, click_index_state, input_mask_state, input_points_state, input_labels_state],
+        outputs=[chatbot, state, image_input, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state],
+        show_progress=True, queue=True
+    )
+
 
 
 
@@ -651,3 +690,4 @@ if __name__ == '__main__':
     iface = create_ui()
     iface.queue(concurrency_count=5, api_open=False, max_size=10)
     iface.launch(server_name="0.0.0.0", enable_queue=True)
+