if __name__ == "__main__":
  args = argparse.ArgumentParser()
  args.add_argument("--gradio", action="store_true", default=False)
  args.add_argument("--input_folder", type=str)
  args = args.parse_args()


  if args.gradio:
    demo.launch(server_name="0.0.0.0", server_port=8080)
  else:
    files = [file for file in os.listdir(args.input_folder) if file.endswith(".jpg") or file.endswith(".jpeg") or file.endswith(".png")]
    for file in tqdm.tqdm(files):
      filename_for_caption = file.split(".")[0]
      caption_file = f"{args.input_folder}/{filename_for_caption}.txt"

      if os.path.exists(caption_file):
        continue

      image = Image.open(f"{args.input_folder}/{file}")
      caption = stream_chat(image.convert("RGB"))
      with open(caption_file, "w", encoding="utf-8") as f:
        f.write(caption)