Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,21 @@
+# 忽略环境文件夹
+myenv/
+__pycache__/
+models/
+releases/
+huggingface/
+
+# 忽略特定的配置文件和图片
+api_settings.json
+saved_prompts.csv
+install_temp.txt
+
+# 忽略系统特定的文件
+Thumbs.db
+.DS_Store
+.vs/
+*.pyproj
+*.user
+*.sln
+
+flash_attn-2.4.2+cu118torch2.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

.ipynb_checkpoints/gpt-caption-checkpoint.py

@@ -0,0 +1,604 @@
+import gradio as gr
+import argparse
+import os
+import shutil
+import threading
+
+import concurrent.futures
+from tqdm import tqdm
+
+import subprocess
+import time
+import requests
+import socket
+
+from lib.Img_Processing import process_images_in_folder, run_script
+from lib.Tag_Processor import modify_file_content, process_tags
+from lib.GPT_Prompt import get_prompts_from_csv, save_prompt, delete_prompt
+from lib.Api_Utils import run_openai_api, save_api_details, get_api_details, downloader, installer, save_state, qwen_api_switch
+from lib.Detecter import detecter
+
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+mod_default, saved_api_key, saved_api_url = get_api_details()
+SUPPORTED_IMAGE_FORMATS = ('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif', '.tiff', '.tif')
+
+# 图像打标
+should_stop = threading.Event()
+def stop_batch_processing():
+    should_stop.set()
+    return "Attempting to stop batch processing. Please wait for the current image to finish."
+
+def process_single_image(api_key, prompt, api_url, image_path, quality, timeout, model="gpt-4o"):
+    save_api_details(api_key, api_url)
+    caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+    print(caption)
+    return caption
+
+def process_batch_images(api_key, prompt, api_url, image_dir, file_handling_mode, quality, timeout, model="gpt-4o"):
+    should_stop.clear()
+    save_api_details(api_key, api_url)
+    results = []
+
+    image_files = []
+    for root, dirs, files in os.walk(image_dir):
+        for file in files:
+            if file.lower().endswith(SUPPORTED_IMAGE_FORMATS):
+                image_files.append(os.path.join(root, file))
+
+    def process_image(filename, file_handling_mode):
+        image_path = os.path.join(image_dir, filename)
+        base_filename = os.path.splitext(filename)[0]
+        caption_filename = f"{base_filename}.txt"
+        caption_path = os.path.join(image_dir, caption_filename)
+
+        if file_handling_mode != "skip/跳过" or not os.path.exists(caption_path):
+            caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+
+            if caption.startswith("Error:") or caption.startswith("API error:"):
+                return handle_error(image_path, caption_path, caption_filename, filename)
+            else:
+                modify_file_content(caption_path, caption, file_handling_mode)
+                return filename, caption_path
+        else:
+            return filename, "Skipped because caption file already exists."
+
+    def handle_error(image_path, caption_path, caption_filename, filename):
+        parent_dir = os.path.dirname(image_dir)
+        error_image_dir = os.path.join(parent_dir, "error_images")
+        if not os.path.exists(error_image_dir):
+            os.makedirs(error_image_dir)
+
+        error_image_path = os.path.join(error_image_dir, filename)
+        error_caption_path = os.path.join(error_image_dir, caption_filename)
+
+        try:
+            shutil.move(image_path, error_image_path)
+            if os.path.exists(caption_path):
+                shutil.move(caption_path, error_caption_path)
+            return filename, "Error handled and image with its caption moved to error directory."
+        except Exception as e:
+            return filename, f"An unexpected error occurred while moving {filename} or {caption_filename}: {e}"
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+        futures = {}
+        for filename in image_files:
+            future = executor.submit(process_image, filename, file_handling_mode)
+            futures[future] = filename  # 将 future 和 filename 映射起来
+        progress = tqdm(total=len(futures), desc="Processing images")
+
+        try:
+            for future in concurrent.futures.as_completed(futures):
+                filename = futures[future]
+                if should_stop.is_set():
+                    for f in futures:
+                        f.cancel()
+                    print("Batch processing was stopped by the user.")
+                    break
+                try:
+                    result = future.result()
+                except Exception as e:
+                    result = (filename, f"An exception occurred: {e}")
+                    print(f"An exception occurred while processing {filename}: {e}")
+                results.append(result)
+                progress.update(1)
+        finally:
+            progress.close()
+            executor.shutdown(wait=False)
+
+    print(f"Processing complete. Total images processed: {len(results)}")
+    return results
+
+def handle_file(image_path, target_path, file_handling_mode):
+    try:
+        if file_handling_mode[:4] == "copy":
+            shutil.copy(image_path, target_path)
+        elif file_handling_mode[:4] == "move":
+            shutil.move(image_path, target_path)
+    except Exception as e:
+        print(f"An exception occurred while handling the file {image_path}: {e}")
+        return f"Error handling file {image_path}: {e}"
+    return
+
+def process_batch_watermark_detection(api_key, prompt, api_url, image_dir, detect_file_handling_mode, quality, timeout,
+                                      watermark_dir, model="gpt-4o"):
+    should_stop.clear()
+    save_api_details(api_key, api_url)
+    results = []
+    prompt = 'Is image have watermark'
+
+    image_files = []
+    for root, dirs, files in os.walk(image_dir):
+        for file in files:
+            if file.lower().endswith(SUPPORTED_IMAGE_FORMATS):
+                image_files.append(os.path.join(root, file))
+
+    def process_image(filename, detect_file_handling_mode, watermark_dir):
+        image_path = os.path.join(image_dir, filename)
+        caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+
+        if caption.startswith("Error:") or caption.startswith("API error:"):
+            return "error"
+
+        # EOI是cog迷之误判？
+        if 'Yes,' in caption and '\'EOI\'' not in caption:
+            target_path = os.path.join(watermark_dir, filename)
+            handle_file(filename, watermark_dir, detect_file_handling_mode)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
+        futures = {}
+        for filename in image_files:
+            future = executor.submit(process_image, filename, detect_file_handling_mode, watermark_dir)
+            futures[future] = filename  # 将 future 和 filename 映射起来
+        progress = tqdm(total=len(futures), desc="Processing images")
+
+        try:
+            for future in concurrent.futures.as_completed(futures):
+                filename = futures[future]  # 获取正在处理的文件名
+                if should_stop.is_set():
+                    for f in futures:
+                        f.cancel()
+                    print("Batch processing was stopped by the user.")
+                    break
+                try:
+                    result = future.result()
+                except Exception as e:
+                    result = (filename, f"An exception occurred: {e}")
+                    print(f"An exception occurred while processing {filename}: {e}")
+                results.append(result)
+                progress.update(1)
+        finally:
+            progress.close()
+            executor.shutdown(wait=False)
+
+    results = f"Total checked images: {len(results)}"
+    return results
+
+def classify_images(api_key, api_url, quality, prompt, timeout, detect_file_handling_mode, image_dir, o_dir, *list_r):
+
+    # 初始化
+    should_stop.clear()
+    save_api_details(api_key, api_url)
+    results = []
+
+    # 检查输入
+    if not os.path.exists(image_dir):
+        return "Error: Image directory does not exist. / 错误：图片目录不存在"
+    if not o_dir:
+        o_dir = os.path.join(image_dir, "classify_output")
+    if not os.path.exists(o_dir):
+        os.makedirs(o_dir)
+
+    # 获取图像
+    image_files = []
+    for root, dirs, files in os.walk(image_dir):
+        for file in files:
+            if file.lower().endswith(SUPPORTED_IMAGE_FORMATS):
+                image_files.append(os.path.join(root, file))
+
+    # 转换列表
+    rules = []
+    for i in range(0, len(list_r), 2):
+        rule_type = list_r[i]
+        rule_input = list_r[i + 1]
+        if rule_type and rule_input:
+            rule_type_bool = rule_type == "Involve / 包含"
+            rules.append((rule_type_bool, rule_input))
+    if rules == []:
+        return "Error: All rules are empty. / 错误：未设置规则"
+
+    # 图像处理
+    def process_image(filename, rules, detect_file_handling_mode, image_dir, o_dir, model="gpt-4o"):
+        image_path = os.path.join(image_dir, filename)
+        caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+
+        if caption.startswith("Error:") or caption.startswith("API error:"):
+            return "error"
+
+        matching_rules = []
+        for rule_bool, rule_input in rules:
+            if (rule_bool and rule_input in caption) or (not rule_bool and rule_input not in caption):
+                matching_rules.append(rule_input)
+
+        if matching_rules:
+            folder_name = "-".join(matching_rules)
+            target_folder = os.path.join(o_dir, folder_name)
+            os.makedirs(target_folder, exist_ok=True)
+            handle_file(filename, target_folder, detect_file_handling_mode)
+        elif matching_rules == []:
+            no_match_folder = os.path.join(o_dir, "no_match")
+            os.makedirs(no_match_folder, exist_ok=True)
+            handle_file(filename, no_match_folder, detect_file_handling_mode)
+
+    # 批量处理
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        futures = {}
+        for filename in image_files:
+            future = executor.submit(process_image, filename, rules, detect_file_handling_mode, image_dir, o_dir)
+            futures[future] = filename  # 将 future 和 filename 映射起来
+        progress = tqdm(total=len(futures), desc="Processing images")
+
+        try:
+            for future in concurrent.futures.as_completed(futures):
+                filename = futures[future]  # 获取正在处理的文件名
+
+                if should_stop.is_set():
+                    for f in futures:
+                        f.cancel()
+                    print("Batch processing was stopped by the user.")
+                    break
+
+                try:
+                    result = future.result()
+                except Exception as e:
+                    result = (filename, f"An exception occurred: {e}")
+                    print(f"An exception occurred while processing {filename}: {e}")
+                results.append(result)
+                progress.update(1)
+
+
+        finally:
+            progress.close()
+            executor.shutdown(wait=False)
+
+    results = f"Total checked images: {len(results)}"
+    return results
+
+# api
+def switch_API(api, state):
+    def is_connection():
+        try:
+            socket.create_connection(("127.0.0.1", 8000), timeout=1)
+            print("API has started.")
+            return True
+        except (socket.timeout, ConnectionRefusedError):
+            return False
+    if api[:3] == 'GPT' or api[:4] == "qwen":
+        if is_connection():
+            requests.post(f"http://127.0.0.1:8000/v1/close")
+        key = saved_api_key
+        url = saved_api_url
+        time_out = 100
+        if api[:4] == "qwen" and url.endswith("/v1/services/aigc/multimodal-generation/generation"):
+            mod = qwen_api_switch(api)
+        else:
+            mod = 'GPT4V'
+        s_state = mod
+
+    elif api[:3] == 'Cog' or api[:4] == "moon" or api[:7] == "MiniCPM":
+        if is_connection():
+            if state != api:
+                requests.post(f"http://127.0.0.1:8000/v1/{api}")
+        else:
+            API_command = f'python openai_api.py --mod {api}'
+            subprocess.Popen(API_command,shell=True)
+            while True:
+                if is_connection():
+                    break
+                else:
+                    print("Retrying...")
+                    time.sleep(2)
+
+        key = ""
+        url = "http://127.0.0.1:8000/v1/chat/completions"
+        time_out = 300
+        s_state = api
+
+    return key, url, time_out, s_state
+
+# UI界面
+with gr.Blocks(title="GPT4V captioner") as demo:
+    gr.Markdown("### Image Captioning with GPT-4-Vision API / 使用 GPT-4-Vision API 进行图像打标")
+
+    with gr.Row():
+        api_key_input = gr.Textbox(label="API Key", placeholder="Enter your GPT-4-Vision API Key here", type="password",
+                                   value=saved_api_key)
+        api_url_input = gr.Textbox(label="API URL", value=saved_api_url or "https://api.openai.com/v1/chat/completions",
+                                   placeholder="Enter the GPT-4-Vision API URL here")
+        api_model_input = gr.Textbox(label="API Model", value="gpt-4o", placeholder="Enter the model name here")
+        quality_choices = ["auto", "high", "low"]
+        quality = gr.Dropdown(choices=quality_choices, label="Image Quality / 图片质量", value="auto")
+        timeout_input = gr.Number(label="Timeout (seconds) / 超时时间（秒）", value=10, step=1)
+
+    prompt_input = gr.Textbox(label="Prompt / 打标需求",
+                              value="As an AI image tagging expert, please provide precise tags for these images to enhance CLIP model's understanding of the content. Employ succinct keywords or phrases, steering clear of elaborate sentences and extraneous conjunctions. Prioritize the tags by relevance. Your tags should capture key elements such as the main subject, setting, artistic style, composition, image quality, color tone, filter, and camera specifications, and any other tags crucial for the image. When tagging photos of people, include specific details like gender, nationality, attire, actions, pose, expressions, accessories, makeup, composition type, age, etc. For other image categories, apply appropriate and common descriptive tags as well. Recognize and tag any celebrities, well-known landmark or IPs if clearly featured in the image. Your tags should be accurate, non-duplicative, and within a 20-75 word count range. These tags will use for image re-creation, so the closer the resemblance to the original image, the better the tag quality. Tags should be comma-separated. Exceptional tagging will be rewarded with $10 per image.",
+                              placeholder="Enter a descriptive prompt",
+                              lines=5)
+
+    with gr.Accordion("Prompt Saving / 提示词存档", open=False):
+        def update_textbox(prompt):
+            return prompt
+        saved_pro = get_prompts_from_csv()
+        saved_prompts_dropdown = gr.Dropdown(label="Saved Prompts / 提示词存档", choices=saved_pro, type="value",interactive=True)
+        with gr.Row():
+            save_prompt_button = gr.Button("Save Prompt / 保存提示词")
+            delete_prompt_button = gr.Button("Delete Prompt / 删除提示词")
+            load_prompt_button = gr.Button("Load Prompt / 读取到输入框")
+
+        save_prompt_button.click(save_prompt, inputs=prompt_input,outputs=[saved_prompts_dropdown])
+        delete_prompt_button.click(delete_prompt, inputs=saved_prompts_dropdown, outputs=[saved_prompts_dropdown])
+        load_prompt_button.click(update_textbox, inputs=saved_prompts_dropdown, outputs=prompt_input)
+
+    with gr.Tab("Image Process / 图片处理"):
+
+        with gr.Tab("Image Zip / 图像预压缩"):
+            with gr.Row():
+                folder_path_input = gr.Textbox(
+                    label="Image Folder Path / 图像文件夹路径",
+                    placeholder="Enter the folder path containing images / 输入包含图像的文件夹路径"
+                )
+                process_images_button = gr.Button("Process Images / 压缩图像")
+
+            with gr.Row():
+                # Add a Markdown component to display the warning message
+                gr.Markdown("""
+            ⚠ **Warning / 警告**: This preprocessing process will resize and compress all image files into jpg format with a total pixel count ≤ 1024×1024 while maintaining the original aspect ratio, ensuring that both dimensions are multiples of 32. **Please make sure to backup your original files before processing!** This procedure can reduce the size of the training set, help to speed up the labeling process, and decrease the time taken to cache latents to disk during training.
+
+            本预处理过程将会在保持原图长宽比情况下，把所有图像文件裁剪压缩为总像素≤1024×1024的jpg文件，并且长宽像素均为32的倍数。**请务必在处理前备份源文件！**该过程可以缩小训练集体积，有助于加快打标速度，并缩短训练过程中的Cache latents to disk时间。
+                """)
+
+            with gr.Row():
+                image_processing_output = gr.Textbox(
+                    label="Image Processing Output / 图像处理输出",
+                    lines=3
+                )
+
+            process_images_button.click(process_images_in_folder,
+                inputs=[folder_path_input],
+                outputs=[image_processing_output])
+
+        with gr.Tab("Single Image / 单图处理"):
+            with gr.Row():
+                image_input = gr.Image(type='filepath', label="Upload Image / 上传图片")
+                single_image_output = gr.Textbox(label="Caption Output / 标签输出")
+            with gr.Row():
+                single_image_submit = gr.Button("Caption Single Image / 图片打标", variant='primary')
+
+        with gr.Tab("Batch Image / 多图批处理"):
+            with gr.Row():
+                batch_dir_input = gr.Textbox(label="Batch Directory / 批量目录",
+                                             placeholder="Enter the directory path containing images for batch processing")
+            with gr.Row():
+                batch_process_submit = gr.Button("Batch Process Images / 批量处理图像", variant='primary')
+            with gr.Row():
+                batch_output = gr.Textbox(label="Batch Processing Output / 批量输出")
+                file_handling_mode = gr.Radio(
+                    choices=["overwrite/覆盖", "prepend/前置插入", "append/末尾追加", "skip/跳过"],
+                    value="overwrite/覆盖",
+                    label="If a caption file exists: / 如果已经存在打标文件: "
+                )
+            with gr.Row():
+                stop_button = gr.Button("Stop Batch Processing / 停止批量处理")
+                stop_button.click(stop_batch_processing, inputs=[], outputs=batch_output)
+
+        with gr.Tab("Failed File Screening / 打标失败文件筛查"):
+            folder_input = gr.Textbox(label="Folder Input / 文件夹输入", placeholder="Enter the directory path")
+            keywords_input = gr.Textbox(placeholder="Enter keywords, e.g., sorry,error / 请输入检索关键词，例如：sorry,error",
+                                        label="Keywords (optional) / 检索关键词（可选）")
+            run_button = gr.Button("Run Script / 运行脚本", variant='primary')
+            output_area = gr.Textbox(label="Script Output / 脚本输出")
+
+            run_button.click(fn=run_script, inputs=[folder_input, keywords_input], outputs=output_area)
+
+    with gr.Tab("Extra Function / 额外功能"):
+
+        gr.Markdown("""
+                    以下功能基于CogVLM开发（GPT4未经测试），极力推荐使用CogVLM-vqa以达到最佳效果。\n
+                    This function is developed based on CogVLM (GPT4 not tested), and it is strongly recommended to use CogVLM-vqa for optimal results.
+                    """)
+
+        with gr.Tab("Watermark Detection / 批量水印检测"):
+            with gr.Row():
+                detect_batch_dir_input = gr.Textbox(label="Image Directory / 图片目录",
+                                                    placeholder="Enter the directory path containing images for batch processing")
+            with gr.Row():
+                watermark_dir = gr.Textbox(label="Watermark Detected Image Directory / 检测到水印的图片目录",
+                                           placeholder="Enter the directory path to move/copy detected images")
+                detect_file_handling_mode = gr.Radio(choices=["move/移动", "copy/复制"], value="move/移动",
+                                                     label="If watermark is detected / 如果图片检测到水印 ")
+            with gr.Row():
+                batch_detect_submit = gr.Button("Batch Detect Images / 批量检测图像", variant='primary')
+            with gr.Row():
+                detect_batch_output = gr.Textbox(label="Output / 结果")
+            with gr.Row():
+                detect_stop_button = gr.Button("Stop Batch Processing / 停止批量处理")
+                detect_stop_button.click(stop_batch_processing, inputs=[], outputs=detect_batch_output)
+        with gr.Tab("Tag Polishing / 标签润色"):
+            gr.Markdown("""
+                    使用其他打标器(如WD1.4)对图片进行打标后，在上方prompt中使用“Describe this image in a very detailed manner and refer these prompt tags:{大括号里替换为放置额外tags文件的目录，会自动读取和图片同名txt。比如 D:\ abc\}”\n
+                    After marking the image using other captioner(such as WD1.4), enter the prompt in the “” marks in the prompt box.
+                        “Describe this image in a very detailed manner and refer these prompt tags:
+                        {This is the txt file path for captions, will automatically read the txt file with the same name as the image. For example, D: \ abc\}”
+                    """)
+        with gr.Tab("Image filtering / 图片筛选"):
+            gr.Markdown("""
+                        使用自定义规则筛选图片，将回答中包含或不包含对应词的图片放入对应规则的文件夹中。输出目录默认在源目录下的classify_output文件夹下。\n
+                        Use custom rules to filter images. Place images containing or not containing corresponding words in the corresponding rule folder in the answer. Output Directory default in source directory \classify_output.
+                        """)
+            with gr.Row():
+                classify_output = gr.Textbox(label="Output / 结果")
+                classify_button = gr.Button("Run / 开始", variant='primary')
+                classify_stop_button = gr.Button("Stop Batch Processing / 停止批量处理")
+            with gr.Row():
+                classify_dir = gr.Textbox(label="Input Image Directory / 输入图片目录",placeholder="Enter the directory path")
+                classify_output_dir = gr.Textbox(label="Output Directory / 输出目录", placeholder="Default source directory / 默认源目录")
+                classify_handling_mode = gr.Radio(label="If meets / 如果符合",choices=["move/移动", "copy/复制"], value="move/移动")
+
+            rule_inputs = []
+            for i in range(1,11):
+                with gr.Row():
+                    rule_type = gr.Dropdown(label="Rule / 规则类型", choices=["","Involve / 包含", "Exclude / 不包含"], value="")
+                    rule_input = gr.Textbox(label="Custom / 自定义", placeholder="Enter the words you need to filter / 输入你需要筛选的词")
+                    rule_inputs.extend([rule_type, rule_input])
+
+    def caption_image(api_key, api_url, prompt, image, quality, timeout, model="gpt-4o"):
+        if image:
+            return process_single_image(api_key, prompt, api_url, image, quality, timeout, model)
+
+    def batch_process(api_key, api_url, prompt, batch_dir, file_handling_mode, quality, timeout, model="gpt-4o"):
+        process_batch_images(api_key, prompt, api_url, batch_dir, file_handling_mode, quality, timeout, model)
+        return "Batch processing complete. Captions saved or updated as '.txt' files next to images."
+
+    def batch_detect(api_key, api_url, prompt, batch_dir, detect_file_handling_mode, quality, timeout, watermark_dir, model="gpt-4o"):
+        results = process_batch_watermark_detection(api_key, prompt, api_url, batch_dir, detect_file_handling_mode,
+                                                    quality, timeout,watermark_dir, model)
+        return results
+
+    single_image_submit.click(caption_image,
+                              inputs=[api_key_input, api_url_input, prompt_input, image_input, quality, timeout_input, api_model_input],
+                              outputs=single_image_output)
+    batch_process_submit.click(batch_process,
+                               inputs=[api_key_input, api_url_input, prompt_input, batch_dir_input,
+                                       file_handling_mode, quality, timeout_input, api_model_input],
+                               outputs=batch_output)
+    batch_detect_submit.click(batch_detect,
+                              inputs=[api_key_input, api_url_input, prompt_input, detect_batch_dir_input,
+                                      detect_file_handling_mode, quality, timeout_input, watermark_dir, api_model_input],
+                              outputs=detect_batch_output)
+
+    classify_button.click(classify_images,
+                          inputs=[api_key_input, api_url_input, quality, prompt_input, timeout_input,
+                                  classify_handling_mode, classify_dir, classify_output_dir] + rule_inputs,
+                          outputs=classify_output)
+    classify_stop_button.click(stop_batch_processing,inputs=[],outputs=classify_output)
+
+    with gr.Tab("Tag Manage / 标签处理"):
+
+        with gr.Row():
+            folder_path_input = gr.Textbox(label="Folder Path / 文件夹路径",
+                                           placeholder="Enter folder path / 在此输入文件夹路径")
+            top_n_input = gr.Number(label="Top N Tags / Top N 标签", value=100)
+            translate_tags_input = gr.Radio(label="Translate Tags to Chinese / 翻译标签",
+                                            choices=["GPT-3.5 translation / GPT3.5翻译",
+                                                     "Free translation / 免费翻译",
+                                                     "No translation / 不翻译"],
+                                            value="No translation / 不翻译")
+            process_tags_button = gr.Button("Process Tags / 处理标签", variant='primary')
+            output_message = gr.Textbox(label="Output Message / 输出信息", interactive=False)
+
+        with gr.Row():
+            tags_to_remove_input = gr.Textbox(label="Tags to Remove / 删除标签",
+                                              placeholder="Enter tags to remove, separated by commas / 输入要删除的标签，用逗号分隔",
+                                              lines=3)
+            tags_to_replace_input = gr.Textbox(label="Tags to Replace / 替换标签",
+                                               placeholder="Enter tags to replace in 'old_tag:new_tag' format, separated by commas / 输入要替换的标签，格式为 '旧标签:新标签'，用逗号分隔",
+                                               lines=3)
+            new_tag_input = gr.Textbox(label="Add New Tag / 添加新标签",
+                                       placeholder="Enter a new tag to add / 输入一个新标签以添加", lines=3)
+            insert_position_input = gr.Radio(label="New Tag Insert Position / 新标签插入位置",
+                                             choices=["Start / 开始", "End / 结束", "Random / 随机"],
+                                             value="Start / 开始")
+
+        with gr.Row():
+            wordcloud_output = gr.Image(label="Word Cloud / 词云")
+            tag_counts_output = gr.Dataframe(label="Top Tags / 高频标签",
+                                             headers=["Tag Name", "Frequency", "Chinese Translation"],
+                                             interactive=True)  # 修改 Dataframe 组件以显示三列
+
+        with gr.Row():
+            network_graph_output = gr.Image(label="Network Graph / 网络图")
+
+        process_tags_button.click(process_tags,
+                                  inputs=[folder_path_input, top_n_input, tags_to_remove_input,
+                                        tags_to_replace_input, new_tag_input, insert_position_input,
+                                        translate_tags_input, api_key_input, api_url_input], # 新增翻译复选框
+                                  outputs=[tag_counts_output, wordcloud_output, network_graph_output, output_message])
+
+
+    # API Config
+    with gr.Tab("API Config / API配置"):
+        # 本地模型配置
+        with gr.Accordion("Local Model / 使用本地模型", open=True):
+            with gr.Row():
+                gr.Markdown("""
+            ⚠ **Warning / 警告**: 
+            This is the API configuration page. To use local model, you need to configure environment and download it.
+                            **Moondream** model **size is about 22g+**, and it takes a long time, Please confirm that the disk space is sufficient.Please confirm that your GPU has sufficient graphics memory ***(approximately 6g)*** 
+                            **CogVLM**, you need to configure environment and download it, which is **approximately 35g+** in size and takes a long time ***(really, really long)***. 
+                            After installation and download, the total space occupied is about ***40g+***. Please confirm that the disk space is sufficient.
+                            In addition, in terms of model selection, the vqa model performs better but slower, while the chat model is faster but slightly weaker.
+                            Please confirm that your GPU has sufficient graphics memory ***(approximately 14g ±)*** when using CogVLM
+                        
+            此为API配置页面，使用本地模型需要配置相关环境并下载模型，
+                            ***moondream***模型大小约为**22g+**需要较长时间，请确认磁盘空间充足。显存需求约为6g，请确认自己的显卡有足够的显存。
+                            ***CogVLM***大小约为**35g+**，需要较长时间 **(真的很长)**。安装以及下载完成后，总占用空间约为40g+，请确认磁盘空间充足。
+                            模型选择上，vqa模型效果更好但是更慢，chat模型更快但是效果略弱。使用CogVLM请确认自己的显卡有足够的显存 ***(约14g±)***
+            """)
+            with gr.Row():
+                detecter_output = gr.Textbox(label="Check Env / 环境检测", interactive=False)
+                detect_button = gr.Button("Check / 检查", variant='primary')
+            with gr.Row():
+                models_select = gr.Radio(label="Choose Models / 选择模型", choices=["moondream","vqa", "chat", "minicpm"], value="moondream")
+                acceleration_select = gr.Radio(label="Choose Default Plz / 选择是否国内加速(如果使用国内加速,请关闭魔法上网)", choices=["CN", "default"],
+                                               value="CN")
+                download_button = gr.Button("Download Models / 下载模型", variant='primary')
+                install_button = gr.Button("Install / 安装", variant='primary')
+
+        # API配置
+        mod_list = [
+            "GPT4V",
+            "qwen-vl-plus",
+            "qwen-vl-max",
+            "moondream",
+            "Cog-vqa",
+            "Cog-chat",
+            "MiniCPM"
+            ]
+        with gr.Row():
+            switch_select = gr.Dropdown(label="Choose API / 选择API", choices=mod_list, value="GPT4V")
+            A_state = gr.Textbox(label="API State / API状态", interactive=False, value=mod_default)
+            switch_button = gr.Button("Switch / 切换", variant='primary')
+            set_default = gr.Button("Set as default / 设为默认", variant='primary')
+
+
+        detect_button.click(detecter, outputs=detecter_output)
+        download_button.click(downloader, inputs=[models_select, acceleration_select],
+                              outputs=detecter_output)
+        install_button.click(installer, outputs=detecter_output)
+        switch_button.click(switch_API, inputs=[switch_select, A_state],
+                            outputs=[api_key_input, api_url_input, timeout_input, A_state])
+        set_default.click(save_state, inputs=[switch_select, api_key_input, api_url_input], outputs=A_state)
+
+    
+    gr.Markdown("""
+        ### Developers: [Jiaye](https://civitai.com/user/jiayev1),  [LEOSAM 是只兔狲](https://civitai.com/user/LEOSAM),  [SleeeepyZhou](https://civitai.com/user/SleeeepyZhou),  [Fok](https://civitai.com/user/fok3827),  [gluttony-10](https://github.com/gluttony-10),  [327](https://github.com/327),  [十字鱼](https://space.bilibili.com/893892)  |  Welcome everyone to add more new features to this project.
+                """)
+
+# 启动参数
+def get_args():
+    parser = argparse.ArgumentParser(description='GPT4V-Image-Captioner启动参数')
+    parser.add_argument("--port", type=int, default="8848", help="占用端口，默认8848")
+    parser.add_argument("--listen", action='store_true', help="打开远程连接，默认关闭")
+    parser.add_argument("--share", action='store_true', help="打开gradio共享，默认关闭")
+    parser.add_argument("--no-browser", action='store_true', help="不要自动打开浏览器，默认关闭")
+    return parser.parse_args()
+
+args = get_args()
+
+if __name__ == "__main__":
+    threading.Thread(target=lambda: switch_API(mod_default, 'GPT')).start()
+    demo.launch(
+        server_name="0.0.0.0" if args.listen else None,
+        server_port=args.port,
+        share=args.share,
+        inbrowser=False if args.no_browser else True
+    )

.ipynb_checkpoints/install_linux_mac-checkpoint.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+
+# Define the Python version
+PYTHON_VERSION=3.10.
+
+# Check if Python is installed and the version is as expected
+if ! command -v python3 --version &>/dev/null || ! python3 --version | grep -q "$PYTHON_VERSION"; then
+    echo "Python is not installed or not the expected version. Please install Python $PYTHON_VERSION."
+    exit 1
+fi
+
+echo "Python is installed."
+
+# Ping google to decide if use mirror
+target_url="www.google.com"
+timeout=3000
+ping -c 1 -W $timeout $target_url -w 3 > /dev/null
+
+if [ $? -ne 0 ]; then
+    echo "Use CN"
+    export PIP_DISABLE_PIP_VERSION_CHECK=1
+    export PIP_NO_CACHE_DIR=1
+    export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
+else
+    echo "Use default"
+fi
+
+# Upgrade pip to the latest version
+pip install --upgrade pip
+
+# Install necessary Python libraries
+pip install -r ./install_script/requirements.txt
+
+echo ""
+echo "Install completed, please run Start to open the GUI"
+echo "安装完毕，请运行Start打开GUI"
+echo ""
+read -p "press any key to continue...
+按任意键继续..."

.ipynb_checkpoints/start_linux_mac-checkpoint.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+export HF_HOME="huggingface"
+
+python ./install_script/check_open.py
+
+python gpt-caption.py --share "$@"
+
+read -p "Press any key to continue . . . "

LICENSE

@@ -0,0 +1,674 @@
+                    GNU GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+                            Preamble
+
+  The GNU General Public License is a free, copyleft license for
+software and other kinds of works.
+
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+the GNU General Public License is intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.  We, the Free Software Foundation, use the
+GNU General Public License for most of our software; it applies also to
+any other work released this way by its authors.  You can apply it to
+your programs, too.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  To protect your rights, we need to prevent others from denying you
+these rights or asking you to surrender the rights.  Therefore, you have
+certain responsibilities if you distribute copies of the software, or if
+you modify it: responsibilities to respect the freedom of others.
+
+  For example, if you distribute copies of such a program, whether
+gratis or for a fee, you must pass on to the recipients the same
+freedoms that you received.  You must make sure that they, too, receive
+or can get the source code.  And you must show them these terms so they
+know their rights.
+
+  Developers that use the GNU GPL protect your rights with two steps:
+(1) assert copyright on the software, and (2) offer you this License
+giving you legal permission to copy, distribute and/or modify it.
+
+  For the developers' and authors' protection, the GPL clearly explains
+that there is no warranty for this free software.  For both users' and
+authors' sake, the GPL requires that modified versions be marked as
+changed, so that their problems will not be attributed erroneously to
+authors of previous versions.
+
+  Some devices are designed to deny users access to install or run
+modified versions of the software inside them, although the manufacturer
+can do so.  This is fundamentally incompatible with the aim of
+protecting users' freedom to change the software.  The systematic
+pattern of such abuse occurs in the area of products for individuals to
+use, which is precisely where it is most unacceptable.  Therefore, we
+have designed this version of the GPL to prohibit the practice for those
+products.  If such problems arise substantially in other domains, we
+stand ready to extend this provision to those domains in future versions
+of the GPL, as needed to protect the freedom of users.
+
+  Finally, every program is threatened constantly by software patents.
+States should not allow patents to restrict development and use of
+software on general-purpose computers, but in those that do, we wish to
+avoid the special danger that patents applied to a free program could
+make it effectively proprietary.  To prevent this, the GPL assures that
+patents cannot be used to render the program non-free.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Use with the GNU Affero General Public License.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU Affero General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the special requirements of the GNU Affero General Public License,
+section 13, concerning interaction through a network will apply to the
+combination as such.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU General Public License from time to time.  Such new versions will
+be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    <one line to give the program's name and a brief idea of what it does.>
+    Copyright (C) <year>  <name of author>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If the program does terminal interaction, make it output a short
+notice like this when it starts in an interactive mode:
+
+    <program>  Copyright (C) <year>  <name of author>
+    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
+    This is free software, and you are welcome to redistribute it
+    under certain conditions; type `show c' for details.
+
+The hypothetical commands `show w' and `show c' should show the appropriate
+parts of the General Public License.  Of course, your program's commands
+might be different; for a GUI interface, you would use an "about box".
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU GPL, see
+<https://www.gnu.org/licenses/>.
+
+  The GNU General Public License does not permit incorporating your program
+into proprietary programs.  If your program is a subroutine library, you
+may consider it more useful to permit linking proprietary applications with
+the library.  If this is what you want to do, use the GNU Lesser General
+Public License instead of this License.  But first, please read
+<https://www.gnu.org/licenses/why-not-lgpl.html>.

README-CN.md

@@ -0,0 +1,127 @@
+# GPT4V-Image-Captioner / GPT4V图像打标器
+
+[软件安装&演示视频](https://www.bilibili.com/video/BV1pw411g7X1/?spm_id_from=333.999.0.0&vd_source=22436c5073194cf38787049c34e04e02)
+
+[英文版说明](https://github.com/jiayev/GPT4V-Image-Captioner/blob/main/README.md)
+
+现在我们有SDwebUI插件版本[sd-webui-GPT4V-Image-Captioner](https://github.com/SleeeepyZhou/sd-webui-GPT4V-Image-Captioner)。
+
+这是一款使用 Gradio 构建，可使用GPT-4-vision API、阿里云[通义千问VL](https://modelscope.cn/organization/qwen)、[Moondream](https://github.com/vikhyat/moondream)模型 或 [CogVLM](https://github.com/THUDM/CogVLM)模型进行图像打标的多功能图像处理工具箱。特色功能包括：
+
+- 一键安装及使用
+- 单图反推及批量打标功能
+- 云端 GPT4V 或 Claude 3 及阿里云[通义千问VL](https://modelscope.cn/organization/qwen) & 本地 [CogVLM](https://github.com/THUDM/CogVLM) 或 [Moondream](https://github.com/vikhyat/moondream)双模型可选
+- 可视化标签分析与处理
+- 图像分桶预压缩
+- 关键词筛查及水印图像识别
+- 图像自定义识别分类
+
+开发者: [Jiaye](https://civitai.com/user/jiayev1), [LEOSAM是只兔狲](https://civitai.com/user/LEOSAM), [SleeeepyZhou](https://space.bilibili.com/360375877), [Fok](https://civitai.com/user/fok3827), GPT4。 欢迎有兴趣的朋友加入，对本项目进行进一步的完善改进。
+
+
+![下载](https://github.com/jiayev/GPT4V-Image-Captioner/assets/16369810/90612e2b-aac1-4368-84d6-482bb660f5aa)
+
+要使用Claude 3，只需将API密钥和URL替换为Claude 3的API密钥和URL (/v1/messages)，并将模型名称更改为"claude-3-opus"（或sonnet）。
+
+# 安装和启动指南
+
+### Windows（如自动安装失败，请参考[手动安装说明](#windows-手动安装说明)）
+
+1. 以管理员权限打开命令提示符，并导航到您想要克隆仓库的目录。
+2. 使用以下命令克隆仓库：
+    ```
+    git clone https://github.com/jiayev/GPT4V-Image-Captioner
+    ```
+3. 双击 `install_windows.bat` 运行，并安装所有必要的依赖项。
+4. 安装完成后，您可以通过双击 `start_windows.bat`来在终端中启动GPT4V-Image-Captioner。
+5. 按住ctrl并点击终端中的URL地址（或复制URL地址在浏览器打开），将在默认浏览器中跳转打开Gradio应用界面。
+6. 请在界面最上方输入OpenAI官方或者第三方的GPT-4V API Key与API Url，设置图像地址后，就可以图像打标了。
+
+
+### Linux / macOS
+
+1. 打开终端，并导航到您想要克隆仓库的目录。
+2. 使用以下命令克隆仓库：
+    ```
+    git clone https://github.com/jiayev/GPT4V-Image-Captioner
+    ```
+3. 导航到克隆的目录：
+    ```
+    cd GPT4V-Image-Captioner
+    ```
+4. 使用以下命令使安装脚本和启动脚本变为可执行：
+    ```
+    chmod +x install_linux_mac.sh; chmod +x Start_linux_mac.sh
+    ```
+5. 执行安装脚本：
+    ```
+    ./install_linux_mac.sh
+    ```
+6. 在终端中执行启动脚本来启动GPT4V-Image-Captioner。
+    ```
+    ./start_linux_mac.sh
+    ```
+7. 复制终端中显示的URL地址，在浏览器中打开Gradio应用界面。
+8. 请在界面最上方输入OpenAI官方或者第三方的GPT-4V API Key与API Url，设置图像地址后，就可以图像打标了。
+
+
+### Windows 手动安装说明
+
+1. 按 `Win + R` 打开命令提示符。键入 `cmd` 然后按 `Enter` 。
+
+2. 使用下面的命令克隆仓库至本地：
+    ```
+    git clone https://github.com/jiayev/GPT4V-Image-Captioner
+    ```
+
+3. 克隆完成后，切换到克隆的目录中：
+    ```
+    cd GPT4V-Image-Captioner
+    ```
+
+4. 在安装依赖库之前，在命令提示符中输入以下命令并按 `Enter` 来检查是否电脑已经安装了 Python：
+    ```
+    python --version
+    ```
+   如果未安装，会显示错误信息。请访问 [Python 官方下载页面](https://www.python.org/downloads/) 并按照指示进行安装。
+
+5. 创建一个名为 `myenv` 的虚拟环境以避免污染全局 Python 环境：
+    ```
+    python -m venv myenv
+    ```
+
+6. 激活你刚创建的虚拟环境：
+    ```
+    myenv\Scripts\activate
+    ```
+
+7. 更新 `pip`至最新版本：
+    ```
+    python -m pip install --upgrade pip
+    ```
+
+8. 在虚拟环境中安装 `requests`、`gradio` 、 `tqdm` 等库：
+    ```
+    pip install scipy networkx wordcloud matplotlib Pillow tqdm gradio requests
+    ```
+
+9. 完成上述步骤后，可通过双击 `Start_windows.bat` 文件来启动 GPT4V-Image-Captioner。
+
+
+## 更新内容
+
+### 2024年1月6日
+- **更智能的一键安装**: 增加了更智能的一键安装 (`install_windows.bat`) 功能，国内的小伙伴不用再看着pip十几kb慢慢爬了，更加国际化(×，简化了程序的安装。
+- **CogVLM支持**: 增加了CogVLM模型的一键安装以及切换页面，没有GPT4的小伙伴也可以靠本地多模态快乐玩耍了（穷哥们狂喜。
+
+### 2024年1月2日
+- **一键安装和一键启动**: 增加了一键安装 (`install_windows.bat` / `install_linux_mac.sh`) 和一键启动 (`Start_windows.bat` / `Start_linux_mac.sh`) 功能，简化了程序的安装和启动过程。
+- **环境说明补充**: 补充了在Windows和Linux环境下程序的安装和启动说明。
+
+### 2024年1月1日
+- **运行加速**: 提高了程序的打标速度。现在可以在2-3秒内完成一张图片的标注。
+- **标签处理**: 对于已有标签的图像文件，提供了以下不同处理选项："覆盖", "前置插入", "结尾追加" 和 "跳过"。
+- **子文件夹处理**: 新程序能够处理文件夹及其子文件夹中的所有图像文件，支持的图像格式包括：'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif', '.tiff', '.tif'。
+- **程序中断**: 增加了在批量打标签过程中中断打标的功能。
+- **报错筛查**: 可以根据关键词，将所有GPT标记失败的图像（例如NSFW内容）移动到新的文件夹中。
+- **本地化**: 增加了对中文的支持。

README.md

@@ -1,12 +1,111 @@
 ---
-title: GPT4V Image Captioner
-emoji: 📉
-colorFrom: green
-colorTo: purple
+title: GPT4V-Image-Captioner
+app_file: gpt-caption.py
 sdk: gradio
-sdk_version: 4.32.0
-app_file: app.py
-pinned: false
+sdk_version: 4.21.0
 ---
+# GPT4V-Image-Captioner / GPT4V图像打标器
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+[中文版说明](https://github.com/jiayev/GPT4V-Image-Captioner/blob/main/README-CN.md)
+
+We now have [sd-webui-GPT4V-Image-Captioner](https://github.com/SleeeepyZhou/sd-webui-GPT4V-Image-Captioner) for SD WebUI
+
+This is a multifunctional image processing toolbox built with Gradio, capable of tagging images using the GPT-4-vision or Claude 3 API, the [cogVLM](https://github.com/THUDM/CogVLM) model, [Qwen-VL](https://huggingface.co/Qwen)(Alibaba Cloud), the [Moondream](https://github.com/vikhyat/moondream) model. 
+
+Key features include:
+
+- One-click installation and use
+- Single image and multi-image batch tagging
+- Choice of online GPT4V or Claude 3 or [Qwen-VL](https://huggingface.co/Qwen)(Alibaba Cloud) & local CogVLM and Moondream models
+- Visual tag analysis and processing
+- Image pre-compression
+- Keyword filtering and watermark image recognition
+
+Developers: [Jiaye](https://civitai.com/user/jiayev1), [LEOSAM是只兔狲](https://civitai.com/user/LEOSAM), [SleeeepyZhou](https://civitai.com/user/SleeeepyZhou), [Fok](https://civitai.com/user/fok3827), GPT4. Welcome everyone to add more new features to this project.
+
+![下载](https://github.com/jiayev/GPT4V-Image-Captioner/assets/16369810/90612e2b-aac1-4368-84d6-482bb660f5aa)
+
+### Please note that the Claude 3 feature is not finished yet.
+To use Claude 3, simply replace the API key and URL with the Claude 3 API key and URL (/v1/messages), and changing the model name to "claude-3-opus" (or sonnet).
+
+# Installation and Startup Guide
+
+### Windows (If the automatic installation fails, please refer to the [Manual Installation Instructions](#windows-manual-installation-instructions))
+
+1. Open Command Prompt as administrator and navigate to the directory where you want to clone the repository.
+2. Clone the repository using the following command:
+    ```
+    git clone https://github.com/jiayev/GPT4V-Image-Captioner
+    ```
+3. Double-click `install_windows.bat` to run and install all necessary dependencies.
+4. After the installation is complete, you can launch the GPT4V-Image-Captioner by double-clicking `start_windows.bat`.
+5. Hold down Ctrl and click on the URL in the terminal (or copy the URL to your browser), which will open the Gradio app interface in your default browser.
+6. Enter the official OpenAI or third-party GPT-4V API Key and API Url at the top of the interface. After setting the image address, you can start tagging the image.
+
+### Linux / macOS
+
+1. Open a terminal and navigate to the directory where you want to clone the repository.
+2. Clone the repository using the following command:
+    ```
+    git clone https://github.com/jiayev/GPT4V-Image-Captioner
+    ```
+3. Navigate to the cloned directory:
+    ```
+    cd GPT4V-Image-Captioner
+    ```
+4. Make the install and start scripts executable with the following command:
+    ```
+    chmod +x install_linux_mac.sh; chmod +x start_linux_mac.sh
+    ```
+5. Execute the install script:
+    ```
+    ./install_linux_mac.sh
+    ```
+6. Launch the GPT4V-Image-Captioner in the terminal by executing the launch script:
+    ```
+    ./start_linux_mac.sh
+    ```
+7. Copy the URL displayed in the terminal and open it in your browser to access the Gradio app interface.
+8. Enter the official OpenAI or third-party GPT-4V API Key and API Url at the top of the interface. After setting the image address, you can start tagging the image.
+
+### Windows Manual Installation Instructions
+
+1. Open the Command Prompt by pressing `Win + R`, typing `cmd`, and then pressing `Enter`.
+
+2. Clone the repository to your local machine using the following command:
+    ```
+    git clone https://github.com/jiayev/GPT4V-Image-Captioner
+    ```
+
+3. Once cloning is complete, navigate to the cloned directory:
+    ```
+    cd GPT4V-Image-Captioner
+    ```
+
+4. Before installing any dependencies, make sure that Python is installed on your system. Check for Python's presence by typing the following command and pressing `Enter` in the Command Prompt:
+    ```
+    python --version
+    ```
+   If Python is not installed, you will get an error message. In that case, please visit the [Python official download page](https://www.python.org/downloads/) and follow the instructions to install it.
+
+5. Create a virtual environment named `myenv` to avoid contaminating the global Python environment:
+    ```
+    python -m venv myenv
+    ```
+
+6. Activate the virtual environment you just created:
+    ```
+    myenv\Scripts\activate
+    ```
+
+7. Update `pip` to date:
+    ```
+    python -m pip install --upgrade pip
+    ```
+
+8. Install libraries within the virtual environment:
+    ```
+    pip install scipy networkx wordcloud matplotlib Pillow tqdm gradio requests
+    ```
+
+9. After completing the steps above, you can start GPT4V-Image-Captioner by double-clicking the `start_windows.bat` file.

gpt-caption.py

@@ -0,0 +1,604 @@
+import gradio as gr
+import argparse
+import os
+import shutil
+import threading
+
+import concurrent.futures
+from tqdm import tqdm
+
+import subprocess
+import time
+import requests
+import socket
+
+from lib.Img_Processing import process_images_in_folder, run_script
+from lib.Tag_Processor import modify_file_content, process_tags
+from lib.GPT_Prompt import get_prompts_from_csv, save_prompt, delete_prompt
+from lib.Api_Utils import run_openai_api, save_api_details, get_api_details, downloader, installer, save_state, qwen_api_switch
+from lib.Detecter import detecter
+
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+mod_default, saved_api_key, saved_api_url = get_api_details()
+SUPPORTED_IMAGE_FORMATS = ('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif', '.tiff', '.tif')
+
+# 图像打标
+should_stop = threading.Event()
+def stop_batch_processing():
+    should_stop.set()
+    return "Attempting to stop batch processing. Please wait for the current image to finish."
+
+def process_single_image(api_key, prompt, api_url, image_path, quality, timeout, model="gpt-4o"):
+    save_api_details(api_key, api_url)
+    caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+    print(caption)
+    return caption
+
+def process_batch_images(api_key, prompt, api_url, image_dir, file_handling_mode, quality, timeout, model="gpt-4o"):
+    should_stop.clear()
+    save_api_details(api_key, api_url)
+    results = []
+
+    image_files = []
+    for root, dirs, files in os.walk(image_dir):
+        for file in files:
+            if file.lower().endswith(SUPPORTED_IMAGE_FORMATS):
+                image_files.append(os.path.join(root, file))
+
+    def process_image(filename, file_handling_mode):
+        image_path = os.path.join(image_dir, filename)
+        base_filename = os.path.splitext(filename)[0]
+        caption_filename = f"{base_filename}.txt"
+        caption_path = os.path.join(image_dir, caption_filename)
+
+        if file_handling_mode != "skip/跳过" or not os.path.exists(caption_path):
+            caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+
+            if caption.startswith("Error:") or caption.startswith("API error:"):
+                return handle_error(image_path, caption_path, caption_filename, filename)
+            else:
+                modify_file_content(caption_path, caption, file_handling_mode)
+                return filename, caption_path
+        else:
+            return filename, "Skipped because caption file already exists."
+
+    def handle_error(image_path, caption_path, caption_filename, filename):
+        parent_dir = os.path.dirname(image_dir)
+        error_image_dir = os.path.join(parent_dir, "error_images")
+        if not os.path.exists(error_image_dir):
+            os.makedirs(error_image_dir)
+
+        error_image_path = os.path.join(error_image_dir, filename)
+        error_caption_path = os.path.join(error_image_dir, caption_filename)
+
+        try:
+            shutil.move(image_path, error_image_path)
+            if os.path.exists(caption_path):
+                shutil.move(caption_path, error_caption_path)
+            return filename, "Error handled and image with its caption moved to error directory."
+        except Exception as e:
+            return filename, f"An unexpected error occurred while moving {filename} or {caption_filename}: {e}"
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+        futures = {}
+        for filename in image_files:
+            future = executor.submit(process_image, filename, file_handling_mode)
+            futures[future] = filename  # 将 future 和 filename 映射起来
+        progress = tqdm(total=len(futures), desc="Processing images")
+
+        try:
+            for future in concurrent.futures.as_completed(futures):
+                filename = futures[future]
+                if should_stop.is_set():
+                    for f in futures:
+                        f.cancel()
+                    print("Batch processing was stopped by the user.")
+                    break
+                try:
+                    result = future.result()
+                except Exception as e:
+                    result = (filename, f"An exception occurred: {e}")
+                    print(f"An exception occurred while processing {filename}: {e}")
+                results.append(result)
+                progress.update(1)
+        finally:
+            progress.close()
+            executor.shutdown(wait=False)
+
+    print(f"Processing complete. Total images processed: {len(results)}")
+    return results
+
+def handle_file(image_path, target_path, file_handling_mode):
+    try:
+        if file_handling_mode[:4] == "copy":
+            shutil.copy(image_path, target_path)
+        elif file_handling_mode[:4] == "move":
+            shutil.move(image_path, target_path)
+    except Exception as e:
+        print(f"An exception occurred while handling the file {image_path}: {e}")
+        return f"Error handling file {image_path}: {e}"
+    return
+
+def process_batch_watermark_detection(api_key, prompt, api_url, image_dir, detect_file_handling_mode, quality, timeout,
+                                      watermark_dir, model="gpt-4o"):
+    should_stop.clear()
+    save_api_details(api_key, api_url)
+    results = []
+    prompt = 'Is image have watermark'
+
+    image_files = []
+    for root, dirs, files in os.walk(image_dir):
+        for file in files:
+            if file.lower().endswith(SUPPORTED_IMAGE_FORMATS):
+                image_files.append(os.path.join(root, file))
+
+    def process_image(filename, detect_file_handling_mode, watermark_dir):
+        image_path = os.path.join(image_dir, filename)
+        caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+
+        if caption.startswith("Error:") or caption.startswith("API error:"):
+            return "error"
+
+        # EOI是cog迷之误判？
+        if 'Yes,' in caption and '\'EOI\'' not in caption:
+            target_path = os.path.join(watermark_dir, filename)
+            handle_file(filename, watermark_dir, detect_file_handling_mode)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
+        futures = {}
+        for filename in image_files:
+            future = executor.submit(process_image, filename, detect_file_handling_mode, watermark_dir)
+            futures[future] = filename  # 将 future 和 filename 映射起来
+        progress = tqdm(total=len(futures), desc="Processing images")
+
+        try:
+            for future in concurrent.futures.as_completed(futures):
+                filename = futures[future]  # 获取正在处理的文件名
+                if should_stop.is_set():
+                    for f in futures:
+                        f.cancel()
+                    print("Batch processing was stopped by the user.")
+                    break
+                try:
+                    result = future.result()
+                except Exception as e:
+                    result = (filename, f"An exception occurred: {e}")
+                    print(f"An exception occurred while processing {filename}: {e}")
+                results.append(result)
+                progress.update(1)
+        finally:
+            progress.close()
+            executor.shutdown(wait=False)
+
+    results = f"Total checked images: {len(results)}"
+    return results
+
+def classify_images(api_key, api_url, quality, prompt, timeout, detect_file_handling_mode, image_dir, o_dir, *list_r):
+
+    # 初始化
+    should_stop.clear()
+    save_api_details(api_key, api_url)
+    results = []
+
+    # 检查输入
+    if not os.path.exists(image_dir):
+        return "Error: Image directory does not exist. / 错误：图片目录不存在"
+    if not o_dir:
+        o_dir = os.path.join(image_dir, "classify_output")
+    if not os.path.exists(o_dir):
+        os.makedirs(o_dir)
+
+    # 获取图像
+    image_files = []
+    for root, dirs, files in os.walk(image_dir):
+        for file in files:
+            if file.lower().endswith(SUPPORTED_IMAGE_FORMATS):
+                image_files.append(os.path.join(root, file))
+
+    # 转换列表
+    rules = []
+    for i in range(0, len(list_r), 2):
+        rule_type = list_r[i]
+        rule_input = list_r[i + 1]
+        if rule_type and rule_input:
+            rule_type_bool = rule_type == "Involve / 包含"
+            rules.append((rule_type_bool, rule_input))
+    if rules == []:
+        return "Error: All rules are empty. / 错误：未设置规则"
+
+    # 图像处理
+    def process_image(filename, rules, detect_file_handling_mode, image_dir, o_dir, model="gpt-4o"):
+        image_path = os.path.join(image_dir, filename)
+        caption = run_openai_api(image_path, prompt, api_key, api_url, quality, timeout, model)
+
+        if caption.startswith("Error:") or caption.startswith("API error:"):
+            return "error"
+
+        matching_rules = []
+        for rule_bool, rule_input in rules:
+            if (rule_bool and rule_input in caption) or (not rule_bool and rule_input not in caption):
+                matching_rules.append(rule_input)
+
+        if matching_rules:
+            folder_name = "-".join(matching_rules)
+            target_folder = os.path.join(o_dir, folder_name)
+            os.makedirs(target_folder, exist_ok=True)
+            handle_file(filename, target_folder, detect_file_handling_mode)
+        elif matching_rules == []:
+            no_match_folder = os.path.join(o_dir, "no_match")
+            os.makedirs(no_match_folder, exist_ok=True)
+            handle_file(filename, no_match_folder, detect_file_handling_mode)
+
+    # 批量处理
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+        futures = {}
+        for filename in image_files:
+            future = executor.submit(process_image, filename, rules, detect_file_handling_mode, image_dir, o_dir)
+            futures[future] = filename  # 将 future 和 filename 映射起来
+        progress = tqdm(total=len(futures), desc="Processing images")
+
+        try:
+            for future in concurrent.futures.as_completed(futures):
+                filename = futures[future]  # 获取正在处理的文件名
+
+                if should_stop.is_set():
+                    for f in futures:
+                        f.cancel()
+                    print("Batch processing was stopped by the user.")
+                    break
+
+                try:
+                    result = future.result()
+                except Exception as e:
+                    result = (filename, f"An exception occurred: {e}")
+                    print(f"An exception occurred while processing {filename}: {e}")
+                results.append(result)
+                progress.update(1)
+
+
+        finally:
+            progress.close()
+            executor.shutdown(wait=False)
+
+    results = f"Total checked images: {len(results)}"
+    return results
+
+# api
+def switch_API(api, state):
+    def is_connection():
+        try:
+            socket.create_connection(("127.0.0.1", 8000), timeout=1)
+            print("API has started.")
+            return True
+        except (socket.timeout, ConnectionRefusedError):
+            return False
+    if api[:3] == 'GPT' or api[:4] == "qwen":
+        if is_connection():
+            requests.post(f"http://127.0.0.1:8000/v1/close")
+        key = saved_api_key
+        url = saved_api_url
+        time_out = 100
+        if api[:4] == "qwen" and url.endswith("/v1/services/aigc/multimodal-generation/generation"):
+            mod = qwen_api_switch(api)
+        else:
+            mod = 'GPT4V'
+        s_state = mod
+
+    elif api[:3] == 'Cog' or api[:4] == "moon" or api[:7] == "MiniCPM":
+        if is_connection():
+            if state != api:
+                requests.post(f"http://127.0.0.1:8000/v1/{api}")
+        else:
+            API_command = f'python openai_api.py --mod {api}'
+            subprocess.Popen(API_command,shell=True)
+            while True:
+                if is_connection():
+                    break
+                else:
+                    print("Retrying...")
+                    time.sleep(2)
+
+        key = ""
+        url = "http://127.0.0.1:8000/v1/chat/completions"
+        time_out = 300
+        s_state = api
+
+    return key, url, time_out, s_state
+
+# UI界面
+with gr.Blocks(title="GPT4V captioner") as demo:
+    gr.Markdown("### Image Captioning with GPT-4-Vision API / 使用 GPT-4-Vision API 进行图像打标")
+
+    with gr.Row():
+        api_key_input = gr.Textbox(label="API Key", placeholder="Enter your GPT-4-Vision API Key here", type="password",
+                                   value=saved_api_key)
+        api_url_input = gr.Textbox(label="API URL", value=saved_api_url or "https://api.openai.com/v1/chat/completions",
+                                   placeholder="Enter the GPT-4-Vision API URL here")
+        api_model_input = gr.Textbox(label="API Model", value="gpt-4o", placeholder="Enter the model name here")
+        quality_choices = ["auto", "high", "low"]
+        quality = gr.Dropdown(choices=quality_choices, label="Image Quality / 图片质量", value="auto")
+        timeout_input = gr.Number(label="Timeout (seconds) / 超时时间（秒）", value=10, step=1)
+
+    prompt_input = gr.Textbox(label="Prompt / 打标需求",
+                              value="As an AI image tagging expert, please provide precise tags for these images to enhance CLIP model's understanding of the content. Employ succinct keywords or phrases, steering clear of elaborate sentences and extraneous conjunctions. Prioritize the tags by relevance. Your tags should capture key elements such as the main subject, setting, artistic style, composition, image quality, color tone, filter, and camera specifications, and any other tags crucial for the image. When tagging photos of people, include specific details like gender, nationality, attire, actions, pose, expressions, accessories, makeup, composition type, age, etc. For other image categories, apply appropriate and common descriptive tags as well. Recognize and tag any celebrities, well-known landmark or IPs if clearly featured in the image. Your tags should be accurate, non-duplicative, and within a 20-75 word count range. These tags will use for image re-creation, so the closer the resemblance to the original image, the better the tag quality. Tags should be comma-separated. Exceptional tagging will be rewarded with $10 per image.",
+                              placeholder="Enter a descriptive prompt",
+                              lines=5)
+
+    with gr.Accordion("Prompt Saving / 提示词存档", open=False):
+        def update_textbox(prompt):
+            return prompt
+        saved_pro = get_prompts_from_csv()
+        saved_prompts_dropdown = gr.Dropdown(label="Saved Prompts / 提示词存档", choices=saved_pro, type="value",interactive=True)
+        with gr.Row():
+            save_prompt_button = gr.Button("Save Prompt / 保存提示词")
+            delete_prompt_button = gr.Button("Delete Prompt / 删除提示词")
+            load_prompt_button = gr.Button("Load Prompt / 读取到输入框")
+
+        save_prompt_button.click(save_prompt, inputs=prompt_input,outputs=[saved_prompts_dropdown])
+        delete_prompt_button.click(delete_prompt, inputs=saved_prompts_dropdown, outputs=[saved_prompts_dropdown])
+        load_prompt_button.click(update_textbox, inputs=saved_prompts_dropdown, outputs=prompt_input)
+
+    with gr.Tab("Image Process / 图片处理"):
+
+        with gr.Tab("Image Zip / 图像预压缩"):
+            with gr.Row():
+                folder_path_input = gr.Textbox(
+                    label="Image Folder Path / 图像文件夹路径",
+                    placeholder="Enter the folder path containing images / 输入包含图像的文件夹路径"
+                )
+                process_images_button = gr.Button("Process Images / 压缩图像")
+
+            with gr.Row():
+                # Add a Markdown component to display the warning message
+                gr.Markdown("""
+            ⚠ **Warning / 警告**: This preprocessing process will resize and compress all image files into jpg format with a total pixel count ≤ 1024×1024 while maintaining the original aspect ratio, ensuring that both dimensions are multiples of 32. **Please make sure to backup your original files before processing!** This procedure can reduce the size of the training set, help to speed up the labeling process, and decrease the time taken to cache latents to disk during training.
+
+            本预处理过程将会在保持原图长宽比情况下，把所有图像文件裁剪压缩为总像素≤1024×1024的jpg文件，并且长宽像素均为32的倍数。**请务必在处理前备份源文件！**该过程可以缩小训练集体积，有助于加快打标速度，并缩短训练过程中的Cache latents to disk时间。
+                """)
+
+            with gr.Row():
+                image_processing_output = gr.Textbox(
+                    label="Image Processing Output / 图像处理输出",
+                    lines=3
+                )
+
+            process_images_button.click(process_images_in_folder,
+                inputs=[folder_path_input],
+                outputs=[image_processing_output])
+
+        with gr.Tab("Single Image / 单图处理"):
+            with gr.Row():
+                image_input = gr.Image(type='filepath', label="Upload Image / 上传图片")
+                single_image_output = gr.Textbox(label="Caption Output / 标签输出")
+            with gr.Row():
+                single_image_submit = gr.Button("Caption Single Image / 图片打标", variant='primary')
+
+        with gr.Tab("Batch Image / 多图批处理"):
+            with gr.Row():
+                batch_dir_input = gr.Textbox(label="Batch Directory / 批量目录",
+                                             placeholder="Enter the directory path containing images for batch processing")
+            with gr.Row():
+                batch_process_submit = gr.Button("Batch Process Images / 批量处理图像", variant='primary')
+            with gr.Row():
+                batch_output = gr.Textbox(label="Batch Processing Output / 批量输出")
+                file_handling_mode = gr.Radio(
+                    choices=["overwrite/覆盖", "prepend/前置插入", "append/末尾追加", "skip/跳过"],
+                    value="overwrite/覆盖",
+                    label="If a caption file exists: / 如果已经存在打标文件: "
+                )
+            with gr.Row():
+                stop_button = gr.Button("Stop Batch Processing / 停止批量处理")
+                stop_button.click(stop_batch_processing, inputs=[], outputs=batch_output)
+
+        with gr.Tab("Failed File Screening / 打标失败文件筛查"):
+            folder_input = gr.Textbox(label="Folder Input / 文件夹输入", placeholder="Enter the directory path")
+            keywords_input = gr.Textbox(placeholder="Enter keywords, e.g., sorry,error / 请输入检索关键词，例如：sorry,error",
+                                        label="Keywords (optional) / 检索关键词（可选）")
+            run_button = gr.Button("Run Script / 运行脚本", variant='primary')
+            output_area = gr.Textbox(label="Script Output / 脚本输出")
+
+            run_button.click(fn=run_script, inputs=[folder_input, keywords_input], outputs=output_area)
+
+    with gr.Tab("Extra Function / 额外功能"):
+
+        gr.Markdown("""
+                    以下功能基于CogVLM开发（GPT4未经测试），极力推荐使用CogVLM-vqa以达到最佳效果。\n
+                    This function is developed based on CogVLM (GPT4 not tested), and it is strongly recommended to use CogVLM-vqa for optimal results.
+                    """)
+
+        with gr.Tab("Watermark Detection / 批量水印检测"):
+            with gr.Row():
+                detect_batch_dir_input = gr.Textbox(label="Image Directory / 图片目录",
+                                                    placeholder="Enter the directory path containing images for batch processing")
+            with gr.Row():
+                watermark_dir = gr.Textbox(label="Watermark Detected Image Directory / 检测到水印的图片目录",
+                                           placeholder="Enter the directory path to move/copy detected images")
+                detect_file_handling_mode = gr.Radio(choices=["move/移动", "copy/复制"], value="move/移动",
+                                                     label="If watermark is detected / 如果图片检测到水印 ")
+            with gr.Row():
+                batch_detect_submit = gr.Button("Batch Detect Images / 批量检测图像", variant='primary')
+            with gr.Row():
+                detect_batch_output = gr.Textbox(label="Output / 结果")
+            with gr.Row():
+                detect_stop_button = gr.Button("Stop Batch Processing / 停止批量处理")
+                detect_stop_button.click(stop_batch_processing, inputs=[], outputs=detect_batch_output)
+        with gr.Tab("Tag Polishing / 标签润色"):
+            gr.Markdown("""
+                    使用其他打标器(如WD1.4)对图片进行打标后，在上方prompt中使用“Describe this image in a very detailed manner and refer these prompt tags:{大括号里替换为放置额外tags文件的目录，会自动读取和图片同名txt。比如 D:\ abc\}”\n
+                    After marking the image using other captioner(such as WD1.4), enter the prompt in the “” marks in the prompt box.
+                        “Describe this image in a very detailed manner and refer these prompt tags:
+                        {This is the txt file path for captions, will automatically read the txt file with the same name as the image. For example, D: \ abc\}”
+                    """)
+        with gr.Tab("Image filtering / 图片筛选"):
+            gr.Markdown("""
+                        使用自定义规则筛选图片，将回答中包含或不包含对应词的图片放入对应规则的文件夹中。输出目录默认在源目录下的classify_output文件夹下。\n
+                        Use custom rules to filter images. Place images containing or not containing corresponding words in the corresponding rule folder in the answer. Output Directory default in source directory \classify_output.
+                        """)
+            with gr.Row():
+                classify_output = gr.Textbox(label="Output / 结果")
+                classify_button = gr.Button("Run / 开始", variant='primary')
+                classify_stop_button = gr.Button("Stop Batch Processing / 停止批量处理")
+            with gr.Row():
+                classify_dir = gr.Textbox(label="Input Image Directory / 输入图片目录",placeholder="Enter the directory path")
+                classify_output_dir = gr.Textbox(label="Output Directory / 输出目录", placeholder="Default source directory / 默认源目录")
+                classify_handling_mode = gr.Radio(label="If meets / 如果符合",choices=["move/移动", "copy/复制"], value="move/移动")
+
+            rule_inputs = []
+            for i in range(1,11):
+                with gr.Row():
+                    rule_type = gr.Dropdown(label="Rule / 规则类型", choices=["","Involve / 包含", "Exclude / 不包含"], value="")
+                    rule_input = gr.Textbox(label="Custom / 自定义", placeholder="Enter the words you need to filter / 输入你需要筛选的词")
+                    rule_inputs.extend([rule_type, rule_input])
+
+    def caption_image(api_key, api_url, prompt, image, quality, timeout, model="gpt-4o"):
+        if image:
+            return process_single_image(api_key, prompt, api_url, image, quality, timeout, model)
+
+    def batch_process(api_key, api_url, prompt, batch_dir, file_handling_mode, quality, timeout, model="gpt-4o"):
+        process_batch_images(api_key, prompt, api_url, batch_dir, file_handling_mode, quality, timeout, model)
+        return "Batch processing complete. Captions saved or updated as '.txt' files next to images."
+
+    def batch_detect(api_key, api_url, prompt, batch_dir, detect_file_handling_mode, quality, timeout, watermark_dir, model="gpt-4o"):
+        results = process_batch_watermark_detection(api_key, prompt, api_url, batch_dir, detect_file_handling_mode,
+                                                    quality, timeout,watermark_dir, model)
+        return results
+
+    single_image_submit.click(caption_image,
+                              inputs=[api_key_input, api_url_input, prompt_input, image_input, quality, timeout_input, api_model_input],
+                              outputs=single_image_output)
+    batch_process_submit.click(batch_process,
+                               inputs=[api_key_input, api_url_input, prompt_input, batch_dir_input,
+                                       file_handling_mode, quality, timeout_input, api_model_input],
+                               outputs=batch_output)
+    batch_detect_submit.click(batch_detect,
+                              inputs=[api_key_input, api_url_input, prompt_input, detect_batch_dir_input,
+                                      detect_file_handling_mode, quality, timeout_input, watermark_dir, api_model_input],
+                              outputs=detect_batch_output)
+
+    classify_button.click(classify_images,
+                          inputs=[api_key_input, api_url_input, quality, prompt_input, timeout_input,
+                                  classify_handling_mode, classify_dir, classify_output_dir] + rule_inputs,
+                          outputs=classify_output)
+    classify_stop_button.click(stop_batch_processing,inputs=[],outputs=classify_output)
+
+    with gr.Tab("Tag Manage / 标签处理"):
+
+        with gr.Row():
+            folder_path_input = gr.Textbox(label="Folder Path / 文件夹路径",
+                                           placeholder="Enter folder path / 在此输入文件夹路径")
+            top_n_input = gr.Number(label="Top N Tags / Top N 标签", value=100)
+            translate_tags_input = gr.Radio(label="Translate Tags to Chinese / 翻译标签",
+                                            choices=["GPT-3.5 translation / GPT3.5翻译",
+                                                     "Free translation / 免费翻译",
+                                                     "No translation / 不翻译"],
+                                            value="No translation / 不翻译")
+            process_tags_button = gr.Button("Process Tags / 处理标签", variant='primary')
+            output_message = gr.Textbox(label="Output Message / 输出信息", interactive=False)
+
+        with gr.Row():
+            tags_to_remove_input = gr.Textbox(label="Tags to Remove / 删除标签",
+                                              placeholder="Enter tags to remove, separated by commas / 输入要删除的标签，用逗号分隔",
+                                              lines=3)
+            tags_to_replace_input = gr.Textbox(label="Tags to Replace / 替换标签",
+                                               placeholder="Enter tags to replace in 'old_tag:new_tag' format, separated by commas / 输入要替换的标签，格式为 '旧标签:新标签'，用逗号分隔",
+                                               lines=3)
+            new_tag_input = gr.Textbox(label="Add New Tag / 添加新标签",
+                                       placeholder="Enter a new tag to add / 输入一个新标签以添加", lines=3)
+            insert_position_input = gr.Radio(label="New Tag Insert Position / 新标签插入位置",
+                                             choices=["Start / 开始", "End / 结束", "Random / 随机"],
+                                             value="Start / 开始")
+
+        with gr.Row():
+            wordcloud_output = gr.Image(label="Word Cloud / 词云")
+            tag_counts_output = gr.Dataframe(label="Top Tags / 高频标签",
+                                             headers=["Tag Name", "Frequency", "Chinese Translation"],
+                                             interactive=True)  # 修改 Dataframe 组件以显示三列
+
+        with gr.Row():
+            network_graph_output = gr.Image(label="Network Graph / 网络图")
+
+        process_tags_button.click(process_tags,
+                                  inputs=[folder_path_input, top_n_input, tags_to_remove_input,
+                                        tags_to_replace_input, new_tag_input, insert_position_input,
+                                        translate_tags_input, api_key_input, api_url_input], # 新增翻译复选框
+                                  outputs=[tag_counts_output, wordcloud_output, network_graph_output, output_message])
+
+
+    # API Config
+    with gr.Tab("API Config / API配置"):
+        # 本地模型配置
+        with gr.Accordion("Local Model / 使用本地模型", open=True):
+            with gr.Row():
+                gr.Markdown("""
+            ⚠ **Warning / 警告**: 
+            This is the API configuration page. To use local model, you need to configure environment and download it.
+                            **Moondream** model **size is about 22g+**, and it takes a long time, Please confirm that the disk space is sufficient.Please confirm that your GPU has sufficient graphics memory ***(approximately 6g)*** 
+                            **CogVLM**, you need to configure environment and download it, which is **approximately 35g+** in size and takes a long time ***(really, really long)***. 
+                            After installation and download, the total space occupied is about ***40g+***. Please confirm that the disk space is sufficient.
+                            In addition, in terms of model selection, the vqa model performs better but slower, while the chat model is faster but slightly weaker.
+                            Please confirm that your GPU has sufficient graphics memory ***(approximately 14g ±)*** when using CogVLM
+                        
+            此为API配置页面，使用本地模型需要配置相关环境并下载模型，
+                            ***moondream***模型大小约为**22g+**需要较长时间，请确认磁盘空间充足。显存需求约为6g，请确认自己的显卡有足够的显存。
+                            ***CogVLM***大小约为**35g+**，需要较长时间 **(真的很长)**。安装以及下载完成后，总占用空间约为40g+，请确认磁盘空间充足。
+                            模型选择上，vqa模型效果更好但是更慢，chat模型更快但是效果略弱。使用CogVLM请确认自己的显卡有足够的显存 ***(约14g±)***
+            """)
+            with gr.Row():
+                detecter_output = gr.Textbox(label="Check Env / 环境检测", interactive=False)
+                detect_button = gr.Button("Check / 检查", variant='primary')
+            with gr.Row():
+                models_select = gr.Radio(label="Choose Models / 选择模型", choices=["moondream","vqa", "chat", "minicpm"], value="moondream")
+                acceleration_select = gr.Radio(label="Choose Default Plz / 选择是否国内加速(如果使用国内加速,请关闭魔法上网)", choices=["CN", "default"],
+                                               value="CN")
+                download_button = gr.Button("Download Models / 下载模型", variant='primary')
+                install_button = gr.Button("Install / 安装", variant='primary')
+
+        # API配置
+        mod_list = [
+            "GPT4V",
+            "qwen-vl-plus",
+            "qwen-vl-max",
+            "moondream",
+            "Cog-vqa",
+            "Cog-chat",
+            "MiniCPM"
+            ]
+        with gr.Row():
+            switch_select = gr.Dropdown(label="Choose API / 选择API", choices=mod_list, value="GPT4V")
+            A_state = gr.Textbox(label="API State / API状态", interactive=False, value=mod_default)
+            switch_button = gr.Button("Switch / 切换", variant='primary')
+            set_default = gr.Button("Set as default / 设为默认", variant='primary')
+
+
+        detect_button.click(detecter, outputs=detecter_output)
+        download_button.click(downloader, inputs=[models_select, acceleration_select],
+                              outputs=detecter_output)
+        install_button.click(installer, outputs=detecter_output)
+        switch_button.click(switch_API, inputs=[switch_select, A_state],
+                            outputs=[api_key_input, api_url_input, timeout_input, A_state])
+        set_default.click(save_state, inputs=[switch_select, api_key_input, api_url_input], outputs=A_state)
+
+    
+    gr.Markdown("""
+        ### Developers: [Jiaye](https://civitai.com/user/jiayev1),  [LEOSAM 是只兔狲](https://civitai.com/user/LEOSAM),  [SleeeepyZhou](https://civitai.com/user/SleeeepyZhou),  [Fok](https://civitai.com/user/fok3827),  [gluttony-10](https://github.com/gluttony-10),  [327](https://github.com/327),  [十字鱼](https://space.bilibili.com/893892)  |  Welcome everyone to add more new features to this project.
+                """)
+
+# 启动参数
+def get_args():
+    parser = argparse.ArgumentParser(description='GPT4V-Image-Captioner启动参数')
+    parser.add_argument("--port", type=int, default="8848", help="占用端口，默认8848")
+    parser.add_argument("--listen", action='store_true', help="打开远程连接，默认关闭")
+    parser.add_argument("--share", action='store_true', help="打开gradio共享，默认关闭")
+    parser.add_argument("--no-browser", action='store_true', help="不要自动打开浏览器，默认关闭")
+    return parser.parse_args()
+
+args = get_args()
+
+if __name__ == "__main__":
+    threading.Thread(target=lambda: switch_API(mod_default, 'GPT')).start()
+    demo.launch(
+        server_name="0.0.0.0" if args.listen else None,
+        server_port=args.port,
+        share=args.share,
+        inbrowser=False if args.no_browser else True
+    )

install_linux_mac.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+
+# Define the Python version
+PYTHON_VERSION=3.10.
+
+# Check if Python is installed and the version is as expected
+if ! command -v python3 --version &>/dev/null || ! python3 --version | grep -q "$PYTHON_VERSION"; then
+    echo "Python is not installed or not the expected version. Please install Python $PYTHON_VERSION."
+    exit 1
+fi
+
+echo "Python is installed."
+
+# Ping google to decide if use mirror
+target_url="www.google.com"
+timeout=3000
+ping -c 1 -W $timeout $target_url -w 3 > /dev/null
+
+if [ $? -ne 0 ]; then
+    echo "Use CN"
+    export PIP_DISABLE_PIP_VERSION_CHECK=1
+    export PIP_NO_CACHE_DIR=1
+    export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
+else
+    echo "Use default"
+fi
+
+# Upgrade pip to the latest version
+pip install --upgrade pip
+
+# Install necessary Python libraries
+pip install -r ./install_script/requirements.txt
+
+echo ""
+echo "Install completed, please run Start to open the GUI"
+echo "安装完毕，请运行Start打开GUI"
+echo ""
+read -p "press any key to continue...
+按任意键继续..."

install_script/check.txt

@@ -0,0 +1,19 @@
+torch
+torchvision
+bitsandbytes
+deepspeed
+transformers
+spacy
+seaborn
+loguru
+streamlit
+timm
+accelerate
+pydantic
+xformers
+requests
+openai
+fastapi
+httpx
+uvicorn
+dashscope

install_script/check_open.py

@@ -0,0 +1,50 @@
+import importlib
+import subprocess
+
+def install_detection(requir_path):
+    # 读
+    file_path = requir_path
+    requirements = []
+    with open(file_path, 'r') as file:
+        for line in file:
+            requirements.append(line.strip())
+
+    # 查
+    missing_libs = []
+    for libs in requirements:
+        if libs.find("==") == -1: #only fix requirements which contain "==", because I don't know how to take "<=",">="... into account at the same time.
+            check_libs = libs
+        else:
+            check_libs = libs[:libs.index("==")]
+        if check_libs == "Pillow": #import PIL instead of import Pillow
+            check_libs = "PIL"
+        try:
+            importlib.import_module(check_libs)
+        except ImportError:
+            if check_libs == "PIL": #switch back
+                check_libs = "Pillow"
+            missing_libs.append(check_libs)
+
+    return missing_libs
+
+def print_missing(missing_libs):
+    # 返
+    if missing_libs == []:
+        return ""
+    else:
+        return f"Not installed libraries: {', '.join(missing_libs)}"
+
+# 启动检查
+def check_open():
+    require_path = "./install_script/requirements.txt"
+    missings = install_detection(require_path)
+    installed = print_missing(missings)
+    if installed == "":
+        return
+    else:
+        print(installed)
+        for lib in missings:
+            subprocess.check_call(["pip", "install", lib])
+
+if __name__ == "__main__":
+    check_open()

install_script/installcog.bat

@@ -0,0 +1,43 @@
+@echo off
+
+call myenv\Scripts\activate
+
+set HF_HOME=huggingface
+REM ͨ���ٶȼ����������ʹ�þ���
+set "target_url=www.google.com"
+set "timeout=4000"
+
+ping %target_url% -n 1 -w %timeout% >nul
+if %ERRORLEVEL% neq 0 (
+    echo Use CN
+    echo ��װ����
+    set PIP_DISABLE_PIP_VERSION_CHECK=1
+    set PIP_NO_CACHE_DIR=1
+    set PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
+
+    echo ��װ torch...
+    pip install torch==2.2.1+cu121 torchvision==0.17.1+cu121 -f https://mirror.sjtu.edu.cn/pytorch-wheels/torch_stable.html -i https://pypi.tuna.tsinghua.edu.cn/simple
+
+) else (
+    echo Use default
+    echo Installing deps...
+    pip install torch==2.2.1+cu121 torchvision==0.17.1+cu121 -i https://download.pytorch.org/whl/cu121
+)
+
+if %ERRORLEVEL% neq 0 (
+    echo torch install failed / torch ��װʧ�� > install_temp.txt
+    pause >nul
+    exit /b 1
+)
+
+
+pip install ./install_script/deepspeed-0.11.2+8ce7471-py3-none-any.whl
+pip install -U -I --no-deps xformers==0.0.25
+pip install -r ./install_script/require.txt
+if %ERRORLEVEL% neq 0 (
+    echo Deps install failed / ������װʧ�� > install_temp.txt
+    pause >nul
+    exit /b 1
+)
+
+echo Install completed / ��װ��� > install_temp.txt

install_script/installcog.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+source myenv/bin/activate
+
+export HF_HOME="huggingface"
+
+target_url="www.google.com"
+timeout=4000
+ping -c 1 -W $timeout $target_url -w 4 > /dev/null
+
+if [ $? -ne 0 ]; then
+    echo "Use CN"
+    echo "安装依赖"
+
+    export PIP_DISABLE_PIP_VERSION_CHECK=1
+    export PIP_NO_CACHE_DIR=1
+    export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
+
+    echo "安装 torch..."
+    pip install torch==2.2.1+cu121 torchvision==0.17.1+cu121 -f https://mirror.sjtu.edu.cn/pytorch-wheels/torch_stable.html -i https://pypi.tuna.tsinghua.edu.cn/simple
+    if [ $? -ne 0 ]; then
+        echo "torch 安装失败" > install_temp.txt
+        exit 1
+    fi
+
+else
+    echo "Use default"
+    echo "Installing deps..."
+    pip install torch==2.2.1+cu121 torchvision==0.17.1+cu121 --extra-index-url https://download.pytorch.org/whl/cu121
+fi
+
+pip install deepspeed
+pip install -U -I --no-deps xformers==0.0.25
+pip install -r ./install_script/require.txt
+if [ $? -ne 0 ]; then
+    echo "Deps install failed / 依赖安装失败" > install_temp.txt
+    exit 1
+fi
+
+echo "Install completed / 安装完毕" > install_temp.txt

install_script/require.txt

@@ -0,0 +1,19 @@
+SwissArmyTransformer>=0.4.9
+transformers==4.37.2
+spacy>=3.6.0
+seaborn>=0.13.0
+loguru~=0.7.2
+streamlit>=1.29.0
+timm>=0.9.12
+accelerate==0.26.1
+pydantic>=2.5.2
+
+requests
+openai>=1.4.0
+fastapi>=0.105.0
+httpx>=0.25.2
+uvicorn~=0.24.0
+einops==0.7.0
+protobuf==4.25.2
+sentencepiece==0.1.99
+bitsandbytes==0.43.0

install_script/requirements.txt

@@ -0,0 +1,11 @@
+scipy==1.12.0
+networkx
+wordcloud
+matplotlib
+Pillow==10.2.0
+tqdm
+gradio==4.21.0
+requests
+huggingface_hub
+GPUtil
+dashscope

install_windows.bat

@@ -0,0 +1,72 @@
+@echo off
+
+REM ���python��װ
+SET PYTHON_VERSION=3.10.2
+SET PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/%PYTHON_VERSION%/python-%PYTHON_VERSION%-amd64.exe
+
+python --version >NUL 2>&1
+if %ERRORLEVEL% neq 0 (
+    echo Python is not installed. Attempting to install Python %PYTHON_VERSION%...
+    bitsadmin /transfer "PythonInstaller" %PYTHON_INSTALLER_URL% python-installer.exe
+    start /wait python-installer.exe /quiet InstallAllUsers=1 PrependPath=1 Include_test=0
+    del /f python-installer.exe
+    python --version >NUL 2>&1
+    if %ERRORLEVEL% neq 0 (
+        echo Failed to install Python.
+        pause >nul
+        exit /b 1
+    )
+)
+
+echo Python installed.
+
+
+REM ���⻷����ⴴ��
+if not exist "myenv" (
+    echo ���ڴ������⻷��...
+    python -m venv myenv
+    if %ERRORLEVEL% neq 0 (
+        echo �������⻷��ʧ�ܣ����� python �Ƿ�װ����Լ� python �汾�Ƿ�Ϊ64λ�汾��python 3.10����python��Ŀ¼�Ƿ��ڻ�������PATH�ڡ�
+        pause >nul
+        exit /b 1
+    )
+)
+
+call myenv\Scripts\activate
+
+
+REM ͨ���ȸ�����������ʹ�þ���
+set "target_url=www.google.com"
+set "timeout=3000"
+
+ping %target_url% -n 1 -w %timeout% >nul
+if %errorlevel% neq 0 (
+    echo Use CN
+    set PIP_DISABLE_PIP_VERSION_CHECK=1
+    set PIP_NO_CACHE_DIR=1
+    set PIP_INDEX_URL=https://mirror.baidu.com/pypi/simple
+) else (
+    echo Use default
+)
+
+set HF_HOME=huggingface
+
+REM ��װ����
+
+echo Installing deps...
+echo ��װ����
+python -m pip install --upgrade pip
+pip install -r ./install_script/requirements.txt
+if %ERRORLEVEL% neq 0 (
+    echo Deps install failed
+    echo ������װʧ�ܡ�
+    pause >nul
+    exit /b 1
+)
+
+echo.
+echo Install completed, please run Start to open the GUI
+echo ��װ��ϣ�������Start��GUI
+echo.
+
+pause

lib/Api_Utils.py

@@ -0,0 +1,382 @@
+import io
+import os
+import time
+import json
+import base64
+import requests
+import subprocess
+import platform
+from PIL import Image
+from requests.adapters import HTTPAdapter
+import re
+from urllib3.util.retry import Retry
+from huggingface_hub import snapshot_download
+
+API_PATH = 'api_settings.json'
+QWEN_MOD = 'qwen-vl-plus'
+DEFAULT_GPT_MODEL = 'gpt-4o'
+DEFAULT_CLAUDE_MODEL = 'claude-3-sonnet'
+
+# 扩展prompt {} 标记功能，从文件读取额外内容
+def addition_prompt_process(prompt, image_path):
+    # 从image_path分离文件名和扩展名，并更改扩展名为.txt
+    if '{' not in prompt and '}' not in prompt:
+        return prompt
+    file_root, _ = os.path.splitext(image_path)
+    new_file_name = os.path.basename(file_root) + ".txt"
+    # 从prompt中提取目录路径
+    directory_path = prompt[prompt.find('{') + 1: prompt.find('}')]
+    # 拼接新的文件路径
+    full_path = os.path.join(directory_path, new_file_name)
+    # 读取full_path指定的文件内容
+    try:
+        with open(full_path, 'r') as file:
+            file_content = file.read()
+    except Exception as e:
+        return f"Error reading file: {e}"
+
+    new_prompt = prompt.replace('{' + directory_path + '}', file_content)
+    return new_prompt
+
+# 通义千问VL
+def is_ali(api_url):
+    if api_url.endswith("/v1/services/aigc/multimodal-generation/generation"):
+        return True
+    else:
+        return False
+    
+def is_claude(api_url, model):
+    if api_url.endswith("v1/messages") or "claude" in model.lower():
+        return True
+    else:
+        return False
+
+def qwen_api_switch(mod):
+    global QWEN_MOD
+    QWEN_MOD = mod
+    return QWEN_MOD
+
+def qwen_api(image_path, prompt, api_key):
+    print(f"QWEN_MOD: {QWEN_MOD}")
+
+    os.environ['DASHSCOPE_API_KEY'] = api_key
+    from dashscope import MultiModalConversation
+    img = f"file://{image_path}"
+    messages = [{
+        'role': 'system',
+        'content': [
+            {'text': 'You are a helpful assistant.'}
+            ]
+        }, {
+        'role':'user',
+        'content': [
+            {'image': img},
+            {'text': prompt},
+            ]
+        }]
+
+    response = MultiModalConversation.call(model=QWEN_MOD, messages=messages, stream=False, max_length=300)
+    if '"status_code": 400' in response:
+        return f"API error: {response}"
+    if response.get("output") and response["output"].get("choices") and response["output"]["choices"][0].get("message") and response["output"]["choices"][0]["message"].get("content"):
+        if response["output"]["choices"][0]["message"]["content"][0].get("text", False):
+            caption = response["output"]["choices"][0]["message"]["content"][0]["text"]
+        else:
+            box_value = response["output"]["choices"][0]["message"]["content"][0]["box"]
+            text_value = response["output"]["choices"][0]["message"]["content"][1]["text"]
+            b_value = re.search(r'<ref>(.*?)</ref>', box_value).group(1)
+            caption = b_value + text_value
+    else:
+        caption = response
+    return caption
+
+def claude_api(image_path, prompt, api_key, api_url, model, quality=None):
+    print(f"CLAUDE_MODEL: {model}")
+    
+    with open(image_path, "rb") as image_file:
+        # Downscale the image
+        image = Image.open(image_file)
+        width, height = image.size
+        if quality:
+            if quality == "high":
+                target = 1024
+            elif quality == "low":
+                target = 512
+            elif quality == "auto":
+                if width >= 1024 or height >= 1024:
+                    target = 1024
+                else:
+                    target = 512
+        else:
+            target = 1024
+            
+        aspect_ratio = width / height
+
+        # Determine the new dimensions while maintaining the aspect ratio
+        if width > target or height > target:
+            if width > height:
+                new_width = target
+                new_height = int(new_width / aspect_ratio)
+            else:
+                new_height = target
+                new_width = int(new_height * aspect_ratio)
+        else:
+            new_width, new_height = width, height
+
+        # Resize the image
+        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
+        # Use buffer to store image
+        buffer = io.BytesIO()
+        resized_image.save(buffer, format="JPEG")
+        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+
+    # Claude API
+    data = {
+        "model": model,
+        "max_tokens": 300,
+        "messages": [
+            {"role": "user", "content": [
+                    {"type": "image", "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": image_base64
+                        }
+                    },
+                    {"type": "text", "text": prompt}
+                ]  
+            }
+        ]
+    }
+
+    # print(f"data: {data}\n")
+
+    headers = {
+        "Content-Type": "application/json",
+        "x-api-key:": api_key,
+        "anthropic-version": "2023-06-01"
+    }
+
+    # 配置重试策略
+    retries = Retry(total=5,
+                    backoff_factor=1,
+                    status_forcelist=[429, 500, 502, 503, 504],
+                    allowed_methods=["HEAD", "GET", "OPTIONS", "POST"])  # 更新参数名
+    
+    with requests.Session() as s:
+        s.mount('https://', HTTPAdapter(max_retries=retries))
+
+        try:
+            response = s.post(api_url, headers=headers, json=data)
+            response.raise_for_status()
+        # 连接错误回显
+        except requests.exceptions.HTTPError as errh:
+            return f"HTTP Error: {errh}"
+        except requests.exceptions.ConnectionError as errc:
+            return f"Error Connecting: {errc}"
+        except requests.exceptions.Timeout as errt:
+            return f"Timeout Error: {errt}"
+        except requests.exceptions.RequestException as err:
+            return f"OOps: Something Else: {err}"
+        
+    try:
+        response_data = response.json()
+        if 'error' in response_data:
+            return f"API error: {response_data['error']['message']}"
+        caption = response_data['content'][0]['text']
+        return caption
+    except Exception as e:
+        return f"Failed to parse the API response: {e}\n{response.text}"
+    
+
+
+# API使用
+def run_openai_api(image_path, prompt, api_key, api_url, quality=None, timeout=10, model=DEFAULT_GPT_MODEL):
+    prompt = addition_prompt_process(prompt, image_path)
+    # print("prompt{}:",prompt)
+
+    # Qwen-VL
+    if is_ali(api_url):
+        return qwen_api(image_path, prompt, api_key)
+    if is_claude(api_url, model):
+        return claude_api(image_path, prompt, api_key, api_url, model, quality)
+    with open(image_path, "rb") as image_file:
+        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
+
+    # GPT-4V
+    data = {
+        "model": model,
+        "messages": [
+            {
+                "role": "user",
+                "content":
+                [
+                    {"type": "image_url", "image_url":
+                        {"url": f"data:image/jpeg;base64,{image_base64}",
+                        "detail": f"{quality}"}
+                    },
+                    {"type": "text", "text": prompt}
+                ]
+            }
+        ],
+        "max_tokens": 300
+    }
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}"
+    }
+
+    # 配置重试策略
+    retries = Retry(total=5,
+                    backoff_factor=1,
+                    status_forcelist=[429, 500, 502, 503, 504],
+                    allowed_methods=["HEAD", "GET", "OPTIONS", "POST"])  # 更新参数名
+
+    with requests.Session() as s:
+        s.mount('https://', HTTPAdapter(max_retries=retries))
+
+        try:
+            response = s.post(api_url, headers=headers, json=data, timeout=timeout)
+            response.raise_for_status()
+        # 连接错误回显
+        except requests.exceptions.HTTPError as errh:
+            return f"HTTP Error: {errh}"
+        except requests.exceptions.ConnectionError as errc:
+            return f"Error Connecting: {errc}"
+        except requests.exceptions.Timeout as errt:
+            return f"Timeout Error: {errt}"
+        except requests.exceptions.RequestException as err:
+            return f"OOps: Something Else: {err}"
+
+    try:
+        response_data = response.json()
+
+        if 'error' in response_data:
+            return f"API error: {response_data['error']['message']}"
+
+        caption = response_data["choices"][0]["message"]["content"]
+        return caption
+    except Exception as e:
+        return f"Failed to parse the API response: {e}\n{response.text}"
+
+
+# API存档
+def save_api_details(api_key, api_url):
+    if is_ali(api_url):
+        settings = {
+        'model' : QWEN_MOD,
+        'api_key': api_key,
+        'api_url': api_url
+        }
+    else:
+        settings = {
+            'model' : 'GPT',
+            'api_key': api_key,
+            'api_url': api_url
+        }
+    # 不记录空的apikey
+    if api_key != "":
+        with open(API_PATH, 'w', encoding='utf-8') as f:
+            json.dump(settings, f)
+
+def save_state(llm, key, url):
+    if llm[:3] == "GPT" or llm[:4] == "qwen":
+        settings = {
+            'model': llm,
+            'api_key': key,
+            'api_url': url
+        }
+
+    elif llm[:3] == "Cog" or llm[:4] == "moon" or llm[:7] == "MiniCPM":
+        settings = {
+            'model' : llm,
+            'api_key': "",
+            'api_url': "http://127.0.0.1:8000/v1/chat/completions"
+        }
+
+    output = f"Set {llm} as default. / {llm}已设为默认"
+    with open(API_PATH, 'w', encoding='utf-8') as f:
+        json.dump(settings, f)
+    return output
+
+# 读取API设置
+def get_api_details():
+    settings_file = API_PATH
+    if os.path.exists(settings_file):
+        with open(settings_file, 'r') as f:
+            settings = json.load(f)
+        if settings.get('model', '') != '':
+            mod = settings.get('model', '')
+            url = settings.get('api_url', '')
+            if mod[:4] == "qwen":
+                global QWEN_MOD
+                QWEN_MOD = mod
+            else:
+                if is_ali(url):
+                    mod = QWEN_MOD
+            return mod, settings.get('api_key', ''), url
+        else:
+            if settings.get('api_key', '') != '':
+                i_key = settings.get('api_key', '')
+                i_url = settings.get('api_url', '')
+                save_api_details(i_key,i_url)
+                with open(settings_file, 'r') as i:
+                    settings = json.load(i)
+                return settings.get('model', ''), settings.get('api_key', ''), settings.get('api_url', '')
+    return 'GPT', '', ''
+
+
+# 本地模型相关
+def downloader(model_type, acceleration):
+    endpoint = 'https://hf-mirror.com' if acceleration == 'CN' else None
+    if model_type == 'vqa' or model_type == 'chat':
+        snapshot_download(
+            repo_id="lmsys/vicuna-7b-v1.5",
+            allow_patterns=["tokenizer*","special_tokens_map.json"],
+            endpoint=endpoint
+        )
+    if model_type == 'vqa':
+        snapshot_download(
+            repo_id="THUDM/cogagent-vqa-hf",
+            local_dir="./models/cogagent-vqa-hf",
+            max_workers=8,
+            endpoint=endpoint
+        )
+    elif model_type == 'chat':
+        snapshot_download(
+            repo_id="THUDM/cogagent-chat-hf",
+            local_dir="./models/cogagent-chat-hf",
+            max_workers=8,
+            endpoint=endpoint
+        )
+    elif model_type == 'moondream':
+        snapshot_download(
+            repo_id="vikhyatk/moondream1",
+            local_dir="./models/moondream",
+            max_workers=8,
+            endpoint=endpoint
+        )
+    elif model_type == 'minicpm':
+        snapshot_download(
+            repo_id="openbmb/MiniCPM-Llama3-V-2_5",
+            local_dir="./models/MiniCPM-Llama3-V-2_5",
+            max_workers=8,
+            endpoint=endpoint
+        )
+    return f"{model_type} Model download completed. / {model_type}模型下载完成"
+
+def installer():
+    if platform.system() == "Windows":
+        install_command = f'.\install_script\installcog.bat'
+    else:
+        install_command = f'./install_script/installcog.sh'
+        subprocess.Popen(f'chmod +x {install_command}', shell=True)
+        subprocess.Popen('', shell=True) #Use an empty subprocess to refresh permission. If deleted, installcog.sh wouldn't launch properly, with Permission denied error
+    subprocess.Popen(install_command, shell=True)
+
+    while not os.path.exists('install_temp.txt'):
+        time.sleep(2)
+    with open('install_temp.txt', 'r') as file:
+        result_string = file.read()
+    os.remove('install_temp.txt')
+    return result_string

lib/Detecter.py

@@ -0,0 +1,60 @@
+import importlib
+import GPUtil
+
+def check_memory():
+    gpus = GPUtil.getGPUs()
+    for gpu in gpus:
+        if gpu.memoryTotal > 12000:
+            return ""
+        elif gpu.memoryTotal > 6000:
+            return "Only MoonDream can be used. / 仅可使用MoonDream"
+    return "Insufficient GPU graphics memory for use. / 显存过小"
+
+def install_detection(requir_path):
+    # 读
+    file_path = requir_path
+    requirements = []
+    with open(file_path, 'r') as file:
+        for line in file:
+            requirements.append(line.strip())
+
+    # 查
+    missing_libs = []
+    for libs in requirements:
+        try:
+            importlib.import_module(libs)
+        except ImportError:
+            missing_libs.append(libs)
+
+    return missing_libs
+
+def print_missing(missing_libs):
+    # 返
+    if missing_libs == []:
+        return ""
+    else:
+        return f"Not installed libraries: {', '.join(missing_libs)}"
+
+def detecter():
+    gpu_check = check_memory()
+    cog_requir = "./install_script/check.txt"
+    installed = print_missing(install_detection(cog_requir))
+
+    if installed == "":
+        return gpu_check + "All listed libraries are installed. / 本地模型依赖安装无误"
+    else:
+        return gpu_check + installed
+
+
+def is_installed(package):
+    try:
+        dist = importlib.metadata.distribution(package)
+    except importlib.metadata.PackageNotFoundError:
+        try:
+            spec = importlib.util.find_spec(package)
+        except ModuleNotFoundError:
+            return False
+
+        return spec is not None
+
+    return dist is not None

lib/Failed_Tagging_File_Screening.py

@@ -0,0 +1,73 @@
+import argparse
+import os
+import shutil
+
+# List of supported image file extensions
+IMAGE_EXTENSIONS = ['.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif', '.tiff', '.tif']
+
+# Global variable to keep track of the number of moved images
+moved_images_count = 0
+
+# Check and move documents and their associated image files with the same name
+def move_files_with_keywords(source_folder, target_folder, keywords):
+    global moved_images_count
+    # Create the target folder if it doesn't exist
+    if not os.path.exists(target_folder):
+        os.makedirs(target_folder)
+
+    # Iterate through all the files and folders within source_folder
+    for root, dirs, files in os.walk(source_folder):
+        for file in files:
+            if file.endswith('.txt'):
+                file_path = os.path.join(root, file)
+                # Check if the text file contains any of the keywords
+                if has_keywords(file_path, keywords):
+                    # Move the text file
+                    shutil.move(file_path, os.path.join(target_folder, file))
+                    # Move the related image files with the same name
+                    move_related_images(root, file, target_folder)
+
+# Check if the text file contains any of the given keywords
+def has_keywords(file_path, keywords):
+    with open(file_path, 'r', encoding='utf-8') as file:
+        content = file.read()
+        return any(keyword.lower() in content.lower() for keyword in keywords)
+
+# Move the image files with the same name as the text file
+def move_related_images(file_dir, text_file, target_folder):
+    global moved_images_count
+    base_name = os.path.splitext(text_file)[0]
+    for ext in IMAGE_EXTENSIONS:
+        image_file = base_name + ext
+        image_path = os.path.join(file_dir, image_file)
+        if os.path.exists(image_path):
+            # Check if the image already exists in the target folder
+            target_image_path = os.path.join(target_folder, image_file)
+            file_counter = 1
+            # Find a unique file name in the target folder
+            while os.path.exists(target_image_path):
+                # Generate a new file name with a counter
+                new_base_name = f"{base_name}_{file_counter}"
+                target_image_path = os.path.join(target_folder, new_base_name + ext)
+                file_counter += 1
+            # Move the image file to the target folder with the new unique name
+            shutil.move(image_path, target_image_path)
+            moved_images_count += 1  # Increment the count for each moved image
+
+def main(image_path, keywords):
+    # The target folder will be created in the same directory as image_path
+    target_folder = os.path.join(os.path.dirname(image_path), 'moved_files')
+    move_files_with_keywords(image_path, target_folder, keywords)
+    # Display the message with the count of moved images and the target folder path
+    print(f"Operation complete / 操作完成. Total images moved: {moved_images_count}. Moved to folder: {target_folder}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Move documents containing keywords and their associated image files with the same name.")
+    parser.add_argument('--image_path', type=str, help='The path to the folder')
+    parser.add_argument('--keywords', type=str, help='List of keywords, separated by commas', default='error,sorry,content')
+    args = parser.parse_args()
+    
+    # Split the received keyword string into a list
+    keywords = args.keywords.split(',')
+    
+    main(args.image_path, keywords)

lib/GPT_Prompt.py

@@ -0,0 +1,38 @@
+import csv
+import os
+import gradio as gr
+
+# GPT Prompt
+PROMPTS_CSV_PATH = "saved_prompts.csv"
+
+def get_prompts_from_csv():
+    if not os.path.exists(PROMPTS_CSV_PATH):
+        return []
+    with open(PROMPTS_CSV_PATH, 'r', newline='', encoding='utf-8') as file:
+        reader = csv.reader(file)
+        # remove empty rows
+        return [row[0] for row in reader if row]
+            
+def save_prompt(prompt):
+    # Append CSV
+    with open(PROMPTS_CSV_PATH, 'a+', newline='', encoding='utf-8') as file:
+        # Move to start
+        file.seek(0)
+        reader = csv.reader(file)
+        existing_prompts = [row[0] for row in reader]
+        if prompt not in existing_prompts:
+            writer = csv.writer(file)
+            writer.writerow([prompt])
+        # Move to end
+        file.seek(0, os.SEEK_END)
+    return gr.Dropdown(label="Saved Prompts", choices=get_prompts_from_csv(), type="value", interactive=True)
+
+def delete_prompt(prompt):
+    lines = []
+    with open(PROMPTS_CSV_PATH, 'r', newline='', encoding='utf-8') as readFile:
+        reader = csv.reader(readFile)
+        lines = [row for row in reader if row and row[0] != prompt]
+    with open(PROMPTS_CSV_PATH, 'w', newline='', encoding='utf-8') as writeFile:
+        writer = csv.writer(writeFile)
+        writer.writerows(lines)
+    return gr.Dropdown(label="Saved Prompts", choices=get_prompts_from_csv(), type="value", interactive=True)

lib/Img_Processing.py

@@ -0,0 +1,147 @@
+import os
+import subprocess
+import concurrent.futures
+
+from PIL import Image
+from tqdm import tqdm
+from PIL import Image, ExifTags
+
+target_resolutions = [
+    (640, 1632),   # 640 * 1632 = 1044480
+    (704, 1472),   # 704 * 1472 = 1036288
+    (768, 1360),   # 768 * 1360 = 1044480
+    (832, 1248),   # 832 * 1248 = 1038336
+    (896, 1152),
+    (960, 1088),   # 960 * 1088 = 1044480
+    (992, 1056),   # 992 * 1056 = 1047552
+    (1024, 1024),  # 1024 * 1024 = 1048576
+    (1056, 992),   # 1056 * 992 = 1047552
+    (1088, 960),   # 1088 * 960 = 1044480
+    (1152, 896),
+    (1248, 832),   # 1248 * 832 = 1038336
+    (1360, 768),   # 1360 * 768 = 1044480
+    (1472, 704),   # 1472 * 704 = 1036288
+    (1632, 640),   # 1632 * 640 = 1044480
+    # (768, 1360),   # 768 * 1360 = 1044480
+    # (1472, 704),   # 1472 * 704 = 1036288
+    # (1024, 1024),  # 1024 * 1024 = 1048576
+]
+
+# 图像预处理
+def apply_exif_orientation(image):
+    try:
+        for orientation in ExifTags.TAGS.keys():
+            if ExifTags.TAGS[orientation] == 'Orientation':
+                break
+        exif = image._getexif()
+
+        if exif is not None:
+            exif = dict(exif.items())
+            orientation_value = exif.get(orientation)
+
+            if orientation_value == 3:
+                image = image.rotate(180, expand=True)
+            elif orientation_value == 6:
+                image = image.rotate(270, expand=True)
+            elif orientation_value == 8:
+                image = image.rotate(90, expand=True)
+    except (AttributeError, KeyError, IndexError, TypeError):
+        # cases: image don't have getexif
+        pass
+
+    return image
+
+def convert_image_to_jpg(img, img_path):
+    """Convert an Image object to JPG."""
+    # Remove extension from original filename and add .jpg
+    base_name = os.path.splitext(img_path)[0]
+    jpg_path = base_name + '.jpg'
+    
+    # Convert image to RGB if it is RGBA (or any other mode)
+    if img.mode != 'RGB':
+        img = img.convert('RGB')
+    
+    img.save(jpg_path, format='JPEG', quality=100)
+
+def process_image(img_path):
+    try:
+        if img_path.lower().endswith((".jpg", ".png", ".bmp", ".gif", ".tif", ".tiff", ".jpeg", ".webp")):
+            img = Image.open(img_path)
+            img = apply_exif_orientation(img)  # Apply the EXIF orientation
+
+            # Convert to 'RGB' if it is 'RGBA' or any other mode
+            img = img.convert('RGB')
+
+            # 计算原图像的宽高比
+            original_aspect_ratio = img.width / img.height
+
+            # 找到最接近原图像宽高比的目标分辨率
+            target_resolution = min(target_resolutions, key=lambda res: abs(original_aspect_ratio - res[0] / res[1]))
+
+            # 计算新的维度
+            if img.width / target_resolution[0] < img.height / target_resolution[1]:
+                new_width = target_resolution[0]
+                new_height = int(img.height * target_resolution[0] / img.width)
+            else:
+                new_height = target_resolution[1]
+                new_width = int(img.width * target_resolution[1] / img.height)
+
+            # 等比缩放图像
+            img = img.resize((new_width, new_height), Image.LANCZOS)
+
+            # 计算裁剪的区域
+            left = int((img.width - target_resolution[0]) / 2)
+            top = int((img.height - target_resolution[1]) / 2)
+            right = int((img.width + target_resolution[0]) / 2)
+            bottom = int((img.height + target_resolution[1]) / 2)
+
+            # 裁剪图像
+            img = img.crop((left, top, right, bottom))
+
+            # 转换并保存图像为JPG格式
+            convert_image_to_jpg(img, img_path)
+
+    except Exception as e:
+        print(f"Error processing image {img_path}: {e}")
+        return None
+
+def delete_non_jpg_files(folder_path):
+    """Delete all non-jpg image files in a directory, but keep txt files."""
+    for dirpath, dirnames, filenames in os.walk(folder_path):
+        for filename in filenames:
+            if not filename.lower().endswith((".jpg", ".txt")):
+                file_path = os.path.join(dirpath, filename)
+                try:
+                    os.remove(file_path)
+                except Exception as e:
+                    print(f"Error occurred while deleting file : {file_path}. Error : {str(e)}")
+
+def process_images_in_folder(folder_path):
+    """
+    Process all images in the given folder according to the target resolutions,
+    then delete all non-jpg files except for .txt files.
+    """
+    processed_files = []
+
+    file_list = [os.path.join(dirpath, filename)
+                 for dirpath, dirnames, filenames in os.walk(folder_path)
+                 for filename in filenames]
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        list(tqdm(executor.map(process_image, file_list), total=len(file_list)))
+
+    delete_non_jpg_files(folder_path)
+    return f"Processed images in folder: {folder_path}"
+
+# 失败检查
+def run_script(folder_path, keywords):
+    keywords = keywords if keywords else "sorry,error"
+    result = subprocess.run(
+        [
+            'python', './lib/Failed_Tagging_File_Screening.py',
+            '--image_path', folder_path,
+            '--keywords', keywords
+        ],
+        capture_output=True, text=True
+    )
+    return result.stdout if result.stdout else "No Output", result.stderr if result.stderr else "No Error"

lib/Tag_Processor.py

@@ -0,0 +1,231 @@
+import os
+import collections
+import random
+
+import matplotlib.pyplot as plt
+import networkx as nx
+
+from wordcloud import WordCloud
+from itertools import combinations
+from lib import Translator
+
+
+def unique_elements(original, addition):
+    original_list = list(map(str.strip, original.split(',')))
+    addition_list = list(map(str.strip, addition.split(',')))
+    combined_list = []
+    seen = set()
+    for item in original_list + addition_list:
+        if item not in seen and item != '':
+            seen.add(item)
+            combined_list.append(item)
+
+    return ', '.join(combined_list)
+
+def save_path(folder_path,file_name):
+    n_path = os.path.join(folder_path, "Tag_analysis")
+    if not os.path.exists(n_path):
+        try:
+            os.makedirs(n_path)
+        except Exception as e:
+            print(f"Error : {e}")
+    save_path = os.path.join(n_path, file_name)
+    return save_path
+
+def modify_tags_in_folder(folder_path, tags_to_remove, tags_to_replace_dict, new_tag, insert_position):
+    for root, dirs, files in os.walk(folder_path):
+        for file in files:
+            if file.endswith('.txt'):
+                file_path = os.path.join(root, file)
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                    tags = [tag.strip() for tag in content.split(',')]
+                # 删除标签
+                tags = [tag for tag in tags if tag not in tags_to_remove]
+                # 替换标签
+                for old_tag, new_tag_replacement in tags_to_replace_dict.items():
+                    tags = [new_tag_replacement if tag == old_tag else tag for tag in tags]
+                # 添加标签
+                if new_tag and new_tag.strip(): 
+                    if insert_position == 'Start / 开始':
+                        tags.insert(0, new_tag.strip())
+                    elif insert_position == 'End / 结束':
+                        tags.append(new_tag.strip())
+                    elif insert_position == 'Random / 随机':
+                        random_index = random.randrange(len(tags)+1)
+                        tags.insert(random_index, new_tag.strip())
+                
+                # 保存修改后的文件
+                with open(file_path, 'w', encoding='utf-8') as f:
+                    updated_content = ', '.join(tags)
+                    f.write(updated_content)
+                    
+    return "Tags modified successfully."
+
+
+# 词云
+def count_tags_in_folder(folder_path, top_n):
+    tags_counter = collections.Counter()
+    for root, dirs, files in os.walk(folder_path):
+        for file in files:
+            if file.endswith('.txt'):
+                file_path = os.path.join(root, file)
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                    tags = content.split(',')
+                    tags = [tag.strip() for tag in tags]
+                    tags_counter.update(tags)
+
+    sorted_tags = sorted(tags_counter.items(), key=lambda x: x[1], reverse=True)
+    return sorted_tags[:top_n]
+
+def generate_network_graph(folder_path, top_n):
+    G = nx.Graph()
+    tags_cooccurrence = collections.defaultdict(int)
+
+    # 读取文件并计算标签的共现关系
+    for root, dirs, files in os.walk(folder_path):
+        for file in files:
+            if file.endswith('.txt'):
+                file_path = os.path.join(root, file)
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                    tags = list(set(content.split(',')))  # 去重
+                    for tag_pair in combinations(tags, 2):
+                        if tag_pair[0].strip() and tag_pair[1].strip():  # 确保标签不是空的
+                            tags_cooccurrence[tag_pair] += 1
+
+    # 只考虑top n个共现关系
+    top_cooccurrences = sorted(tags_cooccurrence.items(), key=lambda x: x[1], reverse=True)[:top_n]
+
+    # 添加边到图中
+    for (tag1, tag2), weight in top_cooccurrences:
+        G.add_edge(tag1.strip(), tag2.strip(), weight=weight)
+
+    # 设置画布大小
+    plt.figure(figsize=(24, 12))
+    
+    # 创建黑色背景
+    gradio_blue = '#0B0F19'
+    plt.gca().set_facecolor(gradio_blue)
+
+    # 为节点设置大小和颜色
+    degrees = dict(G.degree)
+    node_size = [v * 100 for v in degrees.values()]
+    # 使用更鲜亮的颜色映射
+    node_color = [degrees[n] for n in G.nodes]
+
+    # 为边设置宽度
+    edge_width = [G[u][v]['weight'] / 100 for u, v in G.edges]  # 除以10是为了使边宽度合适
+
+    # 计算节点的布局
+    pos = nx.kamada_kawai_layout(G)
+    # pos = nx.spring_layout(G, k=0.5, iterations=50)
+
+    # 绘制节点，使用Plasma配色方案，以适配黑色背景
+    nx.draw_networkx_nodes(G, pos, node_size=node_size,
+                           node_color=node_color, cmap=plt.cm.plasma, alpha=0.8)
+
+    # 绘制边，使用带有透明度的白色
+    nx.draw_networkx_edges(G, pos, width=edge_width, alpha=0.3, edge_color='w')
+
+    # 绘制标签，设置为白色以突出显示
+    nx.draw_networkx_labels(G, pos, font_size=12,
+                            font_weight='bold', font_color='white',
+                            font_family='sans-serif')
+
+    # 移除坐标轴
+    plt.axis('off')
+    
+    # 保存图像
+    save_network = save_path(folder_path,'tag_network.png')
+    plt.savefig(save_network, format='png', dpi=300, bbox_inches='tight', facecolor=gradio_blue)
+    plt.close()
+    return save_network
+
+def generate_wordcloud(folder_path, top):
+    tag_counts = count_tags_in_folder(folder_path, top)
+    wordcloud = WordCloud(width=1600, height=1200, background_color='white')
+    wordcloud.generate_from_frequencies(dict(tag_counts))
+    plt.figure(figsize=(20, 15))
+    plt.imshow(wordcloud, interpolation='bilinear')
+    plt.axis('off')
+    plt.tight_layout(pad=0)
+    save_wordcloud = save_path(folder_path,'tag_wordcloud.png')
+    plt.savefig(save_wordcloud, format='png')
+    plt.close()
+    return save_wordcloud
+
+# Tag处理
+def modify_file_content(file_path, new_content, mode):
+    if mode == "skip/跳过" and os.path.exists(file_path):
+        print(f"Skip writing, as the file {file_path} already exists.")
+        return
+
+    if mode == "overwrite/覆盖" or not os.path.exists(file_path):
+        with open(file_path, 'w', encoding='utf-8') as file:
+            file.write(new_content)
+        return
+
+    with open(file_path, 'r+', encoding='utf-8') as file:
+        existing_content = file.read()
+        file.seek(0)
+        if mode == "prepend/前置插入":
+            combined_content = unique_elements(new_content, existing_content)
+            file.write(combined_content)
+            file.truncate()
+        elif mode == "append/末尾追加":
+            combined_content = unique_elements(existing_content, new_content)
+            file.write(combined_content)
+            file.truncate()
+        else:
+            raise ValueError("Invalid mode. Must be 'overwrite/覆盖', 'prepend/前置插入', or 'append/末尾追加'.")
+
+def process_tags(folder_path, top_n, tags_to_remove, tags_to_replace, new_tag, insert_position, translate, api_key,
+                 api_url):
+    # 解析删除标签
+    tags_to_remove_list = tags_to_remove.split(',') if tags_to_remove else []
+    tags_to_remove_list = [tag.strip() for tag in tags_to_remove_list]
+
+    # 解析替换标签
+    tags_to_replace_dict = {}
+    if tags_to_replace:
+        try:
+            for pair in tags_to_replace.split(','):
+                old_tag, new_replacement_tag = pair.split(':')
+                tags_to_replace_dict[old_tag.strip()] = new_replacement_tag.strip()
+        except ValueError:
+            return "Error: Tags to replace must be in 'old_tag:new_tag' format separated by commas", None, None
+
+    # 修改文件夹中的标签
+    modify_tags_in_folder(folder_path, tags_to_remove_list, tags_to_replace_dict, new_tag, insert_position)
+
+    # 词云及网格图
+    top = int(top_n)
+    wordcloud_path = generate_wordcloud(folder_path, top)
+    networkgraph_path = generate_network_graph(folder_path, top)
+
+    # 翻译Tag功能
+    def truncate_tag(tag, max_length=30): 
+        # 截断过长标签
+        return (tag[:max_length] + '...') if len(tag) > max_length else tag
+    
+    tag_counts = count_tags_in_folder(folder_path, top)
+    
+    if translate.startswith('GPT-3.5 translation / GPT3.5翻译'):
+        translator = Translator.GPTTranslator(api_key, api_url)
+    elif translate.startswith('Free translation / 免费翻译'):
+        translator = Translator.ChineseTranslator()
+    else:
+        translator = None
+    if translator:
+        tags_to_translate = [tag for tag, _ in tag_counts]
+        translations = Translator.translate_tags(translator, tags_to_translate)
+        # 确保 translations 列表长度与 tag_counts 一致
+        translations.extend(["" for _ in range(len(tag_counts) - len(translations))])
+        tag_counts_with_translation = [(truncate_tag(tag_counts[i][0]), tag_counts[i][1], translations[i]) for i in
+                                       range(len(tag_counts))]
+    else:
+        tag_counts_with_translation = [(truncate_tag(tag), count, "") for tag, count in tag_counts]
+
+    return tag_counts_with_translation, wordcloud_path, networkgraph_path, "Tags processed successfully."

lib/Translator.py

@@ -0,0 +1,79 @@
+import requests
+from urllib3.util.retry import Retry
+from requests.adapters import HTTPAdapter
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+class ChineseTranslator:
+    def __init__(self):
+        self.client = requests.Session()
+
+    def translate(self, text):
+        if not text:
+            return None
+
+        payload = {
+            "appid": "105",
+            "sgid": "en",
+            "sbid": "en",
+            "egid": "zh-CN",
+            "ebid": "zh-CN",
+            "content": text,
+            "type": "2",
+        }
+
+        response = self.client.post("https://translate-api-fykz.xiangtatech.com/translation/webs/index", data=payload)
+        if response.status_code == 200:
+            json_data = response.json()
+            by_value = json_data.get("by", "")
+            if not by_value:
+                return None
+            return by_value
+
+        return None
+
+    def close_session(self):
+        self.client.close()
+        
+class GPTTranslator:
+    def __init__(self, api_key, api_url):
+        self.headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}"
+        }
+
+        self.session = requests.Session()
+        retries = Retry(total=5, backoff_factor=0.1, status_forcelist=[429, 500, 502, 503, 504])
+        self.session.mount('https://', HTTPAdapter(max_retries=retries))
+
+        self.api_url = api_url
+
+    def translate(self, text):
+        data = {
+            "model": "gpt-3.5-turbo",
+            "messages": [
+                {"role": "user", "content": f"你是一个英译中专家，请直接返回'{text}'最有可能的三种中文翻译结果，彼此间语义有所区分，结果以逗号间隔."}
+            ]
+        }
+        response = self.session.post(self.api_url, headers=self.headers, json=data)
+        response_data = response.json()
+
+        if response.status_code == 200 and 'choices' in response_data and 'content' in response_data['choices'][0]['message']:
+            return response_data['choices'][0]['message']['content']
+        else:
+            return f"Error or no translation for tag: {text}"
+
+    def close_session(self):
+        self.session.close()
+        
+def translate_tags(translator, tags):
+    translations = [None] * len(tags)
+    
+    with ThreadPoolExecutor(max_workers=50) as executor:
+        future_to_index = {executor.submit(translator.translate, tag): i for i, tag in enumerate(tags)}
+        for future in as_completed(future_to_index):
+            index = future_to_index[future]
+            translations[index] = future.result()
+            
+    translator.close_session()
+            
+    return translations

moondream/__init__.py

@@ -0,0 +1,2 @@
+from .util import detect_device
+from .moondream import Moondream

moondream/configuration_moondream.py

@@ -0,0 +1,74 @@
+from transformers import PretrainedConfig
+
+from typing import Optional
+import math
+
+
+class PhiConfig(PretrainedConfig):
+    model_type = "phi-msft"
+
+    def __init__(
+        self,
+        vocab_size: int = 51200,
+        n_positions: int = 2048,
+        n_embd: int = 2048,
+        n_layer: int = 24,
+        n_inner: Optional[int] = None,
+        n_head: int = 32,
+        n_head_kv: Optional[int] = None,
+        rotary_dim: Optional[int] = 32,
+        activation_function: Optional[str] = "gelu_new",
+        flash_attn: bool = False,
+        flash_rotary: bool = False,
+        fused_dense: bool = False,
+        attn_pdrop: float = 0.0,
+        embd_pdrop: float = 0.0,
+        resid_pdrop: float = 0.0,
+        layer_norm_epsilon: float = 1e-5,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = False,
+        pad_vocab_size_multiple: int = 64,
+        gradient_checkpointing: bool = False,
+        **kwargs
+    ):
+        pad_vocab_size = (
+            math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple
+        )
+        super().__init__(
+            vocab_size=pad_vocab_size,
+            n_positions=n_positions,
+            n_embd=n_embd,
+            n_layer=n_layer,
+            n_inner=n_inner,
+            n_head=n_head,
+            n_head_kv=n_head_kv,
+            activation_function=activation_function,
+            attn_pdrop=attn_pdrop,
+            embd_pdrop=embd_pdrop,
+            resid_pdrop=resid_pdrop,
+            layer_norm_epsilon=layer_norm_epsilon,
+            initializer_range=initializer_range,
+            pad_vocab_size_multiple=pad_vocab_size_multiple,
+            tie_word_embeddings=tie_word_embeddings,
+            gradient_checkpointing=gradient_checkpointing,
+            **kwargs
+        )
+        self.rotary_dim = min(rotary_dim, n_embd // n_head)
+        self.flash_attn = flash_attn
+        self.flash_rotary = flash_rotary
+        self.fused_dense = fused_dense
+
+    attribute_map = {
+        "max_position_embeddings": "n_positions",
+        "hidden_size": "n_embd",
+        "num_attention_heads": "n_head",
+        "num_hidden_layers": "n_layer",
+    }
+
+
+class MoondreamConfig(PretrainedConfig):
+    model_type = "moondream1"
+
+    def __init__(self, **kwargs):
+        self.phi_config = PhiConfig(**kwargs)
+        super().__init__(**kwargs)

moondream/modeling_phi.py

@@ -0,0 +1,720 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+#
+# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
+# Licensed under the BSD 3-Clause License.
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Union, Tuple
+
+import math
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_moondream import PhiConfig
+
+FusedDense = None
+
+
+@dataclass
+class InferenceParams:
+    max_seqlen: int
+    max_batch_size: int
+    seqlen_offset: int = 0
+    batch_size_offset: int = 0
+    key_value_memory_dict: Dict[str, Any] = field(default_factory=dict)
+    lengths_per_sample: torch.Tensor = None
+
+
+class Embedding(nn.Module):
+    def __init__(self, config: PretrainedConfig):
+        super().__init__()
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        self.drop = nn.Dropout(config.embd_pdrop)
+
+    def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        return self.drop(self.wte(input_ids.view(-1, input_ids.size(-1))))
+
+
+def _apply_rotary_emb(x, cos, sin):
+    seqlen, rotary_dim = x.size(1), cos.size(1) * 2
+    x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:]
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    c, s = cos[:seqlen].unsqueeze(1), sin[:seqlen].unsqueeze(1)
+    x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1)
+    return torch.cat([x_rot.to(x.dtype), x_pass], dim=-1)
+
+
+def _apply_rotary_emb_kv(
+    kv: torch.FloatTensor, cos: torch.FloatTensor, sin: torch.FloatTensor
+) -> torch.FloatTensor:
+    seqlen, rotary_dim = kv.shape[1], cos.shape[-1] * 2
+    k_rot = kv[:, :, 0, :, :rotary_dim].chunk(2, dim=-1)
+    k_pass = kv[:, :, 0, :, rotary_dim:]
+    c, s = cos[:seqlen].unsqueeze(1), sin[:seqlen].unsqueeze(1)
+    k_rot = torch.cat(
+        [k_rot[0] * c - k_rot[1] * s, k_rot[0] * s + k_rot[1] * c], dim=-1
+    )
+    return torch.cat(
+        [torch.cat([k_rot, k_pass], dim=-1).unsqueeze(2), kv[:, :, 1:2, :, :]], dim=2
+    )
+
+
+def _apply_rotary_emb_qkv(
+    qkv: torch.FloatTensor, cos: torch.FloatTensor, sin: torch.FloatTensor
+) -> torch.FloatTensor:
+    seqlen, rotary_dim = qkv.shape[1], cos.shape[1] * 2
+
+    c = cos[:seqlen].unsqueeze(1)
+    s = sin[:seqlen].unsqueeze(1)
+
+    qkv_rot = torch.stack(
+        [
+            torch.cat(
+                [
+                    qkv[:, :, i, :, : rotary_dim // 2] * c
+                    - qkv[:, :, i, :, rotary_dim // 2 : rotary_dim] * s,
+                    qkv[:, :, i, :, : rotary_dim // 2] * s
+                    + qkv[:, :, i, :, rotary_dim // 2 : rotary_dim] * c,
+                ],
+                dim=-1,
+            ).to(qkv.dtype)
+            for i in range(2)
+        ],
+        dim=2,
+    )
+
+    qkv_pass = qkv[:, :, :2, :, rotary_dim:].unsqueeze(2)
+    qkv_v = qkv[:, :, 2:3, :, :]
+    return torch.cat([qkv_rot, qkv_pass, qkv_v], dim=2)
+
+
+class RotaryEmbedding(nn.Module):
+    # Enhanced Transformer with Rotary Position Embedding (https://arxiv.org/pdf/2104.09864.pdf)
+    def __init__(
+        self,
+        dim: int,
+        base: int = 10000,
+        scale_base: Optional[float] = None,
+        pos_idx_in_fp32: bool = True,
+        max_position_embeddings: int = 2048,
+        device: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        # fp32 is preferred since the output of `torch.arange` can be quite large and bf16 would lose a lot of precision
+        self.dim, self.base, self.pos_idx_in_fp32, self.device = (
+            dim,
+            float(base),
+            pos_idx_in_fp32,
+            device,
+        )
+        self.max_position_embeddings = max_position_embeddings
+        if scale_base is not None:
+            raise NotImplementedError
+
+        # Generate and register the non-trainable buffers
+        self.register_buffer(
+            "inv_freq", self._compute_inv_freq(device), persistent=False
+        )
+        self.register_buffer(
+            "scale", self._calculate_scale(dim, scale_base, device), persistent=False
+        )
+        self._update_cos_sin_cache(
+            max_position_embeddings, device=device, dtype=torch.float32
+        )
+
+    def _calculate_scale(self, dim, scale_base, device):
+        return (
+            (
+                (
+                    torch.arange(0, dim, 2, device=device, dtype=torch.float32)
+                    + 0.4 * dim
+                )
+                / (1.4 * dim)
+            )
+            if scale_base is not None
+            else None
+        )
+
+    def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:
+        return 1.0 / (
+            self.base
+            ** (
+                torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)
+                / self.dim
+            )
+        )
+
+    def _update_cos_sin_cache(
+        self,
+        seqlen: int,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        self._seq_len_cached = seqlen
+        t = torch.arange(
+            seqlen,
+            device=device,
+            dtype=torch.float32 if self.pos_idx_in_fp32 else self.inv_freq.dtype,
+        )
+        inv_freq = (
+            self._compute_inv_freq(device=device)
+            if self.pos_idx_in_fp32 and self.inv_freq.dtype != torch.float32
+            else self.inv_freq
+        )
+
+        freqs = torch.outer(t, inv_freq)
+
+        def apply_scale(freqs, scale, operator, dtype):
+            result = operator(freqs)
+            return (result / scale).to(dtype) if scale is not None else result.to(dtype)
+
+        if scale := self.scale:
+            power = (
+                torch.arange(seqlen, dtype=scale.dtype, device=scale.device)
+                - seqlen // 2
+            ) / self.scale_base
+            scale = scale.to(device=power.device) ** power.unsqueeze(1)
+
+        self._cos_cached = apply_scale(
+            freqs, 1 / scale if scale is not None else None, torch.cos, dtype
+        )
+        self._sin_cached = apply_scale(
+            freqs, 1 / scale if scale is not None else None, torch.sin, dtype
+        )
+        if scale is not None:
+            self._cos_k_cached = apply_scale(freqs, scale, torch.cos, dtype)
+            self._sin_k_cached = apply_scale(freqs, scale, torch.sin, dtype)
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        kv: Optional[torch.Tensor] = None,
+        seqlen_offset: int = 0,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        should_update = (
+            self._seq_len_cached < qkv.shape[1] + seqlen_offset
+            or self._cos_cached.device != qkv.device
+            or self._cos_cached.dtype != qkv.dtype
+            or (self.training and self._cos_cached.is_inference())
+        )
+
+        if should_update:
+            self._update_cos_sin_cache(
+                qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype
+            )
+
+        offset_cos = self._cos_cached[seqlen_offset:]
+        offset_sin = self._sin_cached[seqlen_offset:]
+
+        if kv is None:
+            return _apply_rotary_emb_qkv(qkv, offset_cos, offset_sin)
+        else:
+            return _apply_rotary_emb(qkv, offset_cos, offset_sin), _apply_rotary_emb_kv(
+                kv, offset_cos, offset_sin
+            )
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        n_inner: Optional[int] = None,
+        act_fn: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        n_inner = n_inner or getattr(config, "n_inner", None) or 4 * config.n_embd
+        act_fn = act_fn or config.activation_function
+
+        self.fc1 = nn.Linear(config.n_embd, n_inner)
+        self.fc2 = nn.Linear(n_inner, config.n_embd)
+        self.act = ACT2FN[act_fn]
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        return self.fc2(self.act(self.fc1(hidden_states)))
+
+
+# Flash Attention (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py)
+class SelfAttention(nn.Module):
+    def __init__(
+        self,
+        causal: bool = True,
+        softmax_scale: Optional[float] = None,
+        attention_dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.drop = nn.Dropout(attention_dropout)
+
+    @torch.autocast("cpu", enabled=False)
+    @torch.autocast("cuda", enabled=False)
+    def forward(
+        self,
+        qkv: torch.FloatTensor,
+        causal: Optional[bool] = None,
+        key_padding_mask: Optional[torch.BoolTensor] = None,
+    ):
+        q, k, v = qkv.chunk(3, dim=-1)
+        scale = self.softmax_scale or 1.0 / q.size(-1) ** 0.5
+
+        scores = (
+            torch.einsum("bthd,bshd->bhts", q.to(torch.float32), k.to(torch.float32))
+            * scale
+        )
+        if causal or self.causal:
+            scores.triu_(1).fill_(-10000.0)
+        if key_padding_mask is not None:
+            scores.masked_fill_(key_padding_mask[:, None, None, :], -10000.0)
+
+        attn = self.drop(torch.softmax(scores, dim=-1).to(v.dtype))
+        return torch.einsum("bhts,bshd->bthd", attn, v)
+
+
+# Flash Attention (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py)
+class CrossAttention(nn.Module):
+    def __init__(self, causal=True, softmax_scale=None, attention_dropout=0.0):
+        super().__init__()
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.drop = nn.Dropout(attention_dropout)
+
+    @torch.autocast("cpu", enabled=False)
+    @torch.autocast("cuda", enabled=False)
+    def forward(
+        self,
+        q: torch.FloatTensor,
+        kv: torch.FloatTensor,
+        causal: bool = None,
+        key_padding_mask: Optional[torch.BoolTensor] = None,
+    ) -> torch.FloatTensor:
+        batch_size, seqlen_q = q.shape[0], q.shape[1]
+        seqlen_k = kv.shape[1]
+
+        if kv.shape[3] != q.shape[2]:
+            kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])
+        k, v = kv.unbind(dim=2)
+
+        q = q.to(torch.float32)
+        k = k.to(torch.float32)
+
+        causal = self.causal if causal is None else causal
+        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+
+        # Autocast is manually disabled to avoid `torch.einsum` performing the operation using float16, which might lead to overflow
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+        if key_padding_mask is not None:
+            padding_mask = torch.full(
+                (batch_size, seqlen_k),
+                -10000.0,
+                dtype=scores.dtype,
+                device=scores.device,
+            )
+            padding_mask.masked_fill_(key_padding_mask, 0.0)
+            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
+
+        if causal:
+            rows = rearrange(
+                torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1"
+            )
+            cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)
+            causal_mask = cols > rows + seqlen_k - seqlen_q
+            scores = scores.masked_fill(causal_mask, -10000.0)
+
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+        attention = self.drop(attention)
+        output = torch.einsum("bhts,bshd->bthd", attention, v)
+
+        return output
+
+
+def _find_mha_dims(
+    config: PretrainedConfig,
+    n_head: Optional[int] = None,
+    n_head_kv: Optional[int] = None,
+    head_dim: Optional[int] = None,
+) -> Tuple[int, int]:
+    if n_head is None and head_dim is None:
+        head_dim = config.n_embd // config.n_head
+        n_head = config.n_head
+    elif n_head is None or head_dim is None:
+        raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
+    if n_head_kv is None:
+        n_head_kv = getattr(config, "n_head_kv", None) or n_head
+    return n_head, n_head_kv, head_dim
+
+
+def _update_kv_cache(
+    kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int
+) -> torch.FloatTensor:
+    num_heads, head_dim = kv.shape[-2:]
+    layer_memory = inference_params.key_value_memory_dict.setdefault(
+        layer_idx,
+        torch.empty(
+            inference_params.max_batch_size,
+            inference_params.max_seqlen,
+            2,
+            num_heads,
+            head_dim,
+            dtype=kv.dtype,
+            device=kv.device,
+        ),
+    )
+
+    batch_slice = slice(
+        inference_params.batch_size_offset,
+        inference_params.batch_size_offset + kv.shape[0],
+    )
+    seqlen_slice = slice(
+        inference_params.seqlen_offset, inference_params.seqlen_offset + kv.shape[1]
+    )
+
+    if seqlen_slice.stop >= inference_params.max_seqlen:
+        layer_memory = torch.cat((layer_memory, kv), dim=1)
+        inference_params.key_value_memory_dict[layer_idx] = layer_memory
+
+    layer_memory[batch_slice, seqlen_slice, ...] = kv
+    return layer_memory[batch_slice, : seqlen_slice.stop, ...]
+
+
+# Multi-head attention layer with rotary embeddings
+class MHA(nn.Module):
+    def __init__(
+        self,
+        config,
+        dtype=None,
+        device=None,
+        rotary_dim=None,
+        rotary_base=10000.0,
+        rotary_scale_base=None,
+        n_head=None,
+        n_head_kv=None,
+        head_dim=None,
+        bias=True,
+        causal=True,
+        softmax_scale=None,
+        layer_idx=None,
+        return_residual=False,
+        checkpointing=False,
+    ):
+        super().__init__()
+
+        # Set rotary embedding if specified
+        self.rotary_dim = rotary_dim or getattr(config, "rotary_dim", 0)
+        if self.rotary_dim:
+            self.rotary_emb = RotaryEmbedding(
+                self.rotary_dim,
+                base=rotary_base,
+                scale_base=rotary_scale_base,
+                device=device,
+                max_position_embeddings=config.n_positions,
+            )
+
+        # Determine MHA dims from arguments or config
+        self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(
+            config, n_head, n_head_kv, head_dim
+        )
+        op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
+        hidden_size = config.n_embd
+
+        # Choose Linear class based on config, FusedDense is optional
+        LinearClass = (
+            FusedDense if config.fused_dense and FusedDense is not None else nn.Linear
+        )
+        self.Wqkv = LinearClass(
+            hidden_size, op_size, bias=bias, device=device, dtype=dtype
+        )
+        self.out_proj = LinearClass(
+            hidden_size, hidden_size, bias=bias, device=device, dtype=dtype
+        )
+
+        # Initialize attention mechanisms
+        attn_kwargs = {
+            "causal": causal,
+            "softmax_scale": softmax_scale,
+            "attention_dropout": config.attn_pdrop,
+        }
+        self.inner_attn = SelfAttention(**attn_kwargs)
+        self.inner_cross_attn = CrossAttention(**attn_kwargs)
+
+        self.layer_idx = layer_idx
+        self.return_residual = return_residual
+        self.checkpointing = checkpointing
+
+    def _forward_self_attn(
+        self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]
+    ) -> torch.FloatTensor:
+        qkv = rearrange(
+            self.Wqkv(x), "... (three h d) -> ... three h d", three=3, d=self.head_dim
+        )
+        if self.rotary_dim > 0:
+            qkv = self.rotary_emb(qkv)
+        attn_func = (
+            torch.utils.checkpoint.checkpoint
+            if self.checkpointing
+            else lambda f, *args, **kwargs: f(*args, **kwargs)
+        )
+        return attn_func(self.inner_attn, qkv, key_padding_mask=key_padding_mask)
+
+    def _forward_cross_attn(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams],
+        key_padding_mask: Optional[torch.BoolTensor],
+    ) -> torch.FloatTensor:
+        qkv = self.Wqkv(x)
+        q, kv = (
+            qkv[..., : self.n_head * self.head_dim],
+            qkv[..., self.n_head * self.head_dim :],
+        )
+        q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
+        kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
+
+        seqlen_offset = (
+            past_key_values.seqlen_offset if past_key_values is not None else 0
+        )
+        causal = None if seqlen_offset == 0 else False
+        if self.rotary_dim > 0:
+            q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)
+
+        if past_key_values is not None:
+            kv = _update_kv_cache(kv, past_key_values, self.layer_idx)
+
+        attn_func = (
+            torch.utils.checkpoint.checkpoint
+            if self.checkpointing
+            else lambda fn, *args, **kwargs: fn(*args, **kwargs)
+        )
+
+        return attn_func(
+            self.inner_cross_attn,
+            q,
+            kv,
+            key_padding_mask=key_padding_mask,
+            causal=causal,
+        )
+
+    def forward(
+        self,
+        x: torch.FloatTensor,
+        past_key_values: Optional[InferenceParams] = None,
+        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        attention_mask = attention_mask.bool() if attention_mask is not None else None
+        use_cross_attn = self.n_head != self.n_head_kv or past_key_values is not None
+        attn_output_function = (
+            self._forward_cross_attn if use_cross_attn else self._forward_self_attn
+        )
+        attn_output = (
+            attn_output_function(x, past_key_values, attention_mask)
+            if use_cross_attn
+            else attn_output_function(x, attention_mask)
+        )
+        output = self.out_proj(rearrange(attn_output, "... h d -> ... (h d)"))
+        return (output, x) if self.return_residual else output
+
+
+# Parallel block. This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
+class ParallelBlock(nn.Module):
+    def __init__(self, config: PretrainedConfig, block_idx: Optional[int] = None):
+        super().__init__()
+        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+        self.block_idx = block_idx
+        self.mixer = MHA(config, layer_idx=block_idx)
+        self.mlp = MLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+        hidden_states = self.ln(hidden_states)
+
+        attn_outputs = self.mixer(
+            hidden_states,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+        )
+        if isinstance(attn_outputs, tuple):
+            attn_outputs = attn_outputs[0]
+
+        attn_outputs = self.resid_dropout(attn_outputs)
+        feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
+        return attn_outputs + feed_forward_hidden_states + residual
+
+
+class CausalLMHead(nn.Module):
+    """Causal Language Modeling head. Simplified version."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.linear = nn.Linear(config.n_embd, config.vocab_size)
+
+    def forward(self, hidden_states):
+        return self.linear(self.ln(hidden_states)).to(torch.float32)
+
+
+# Improving Language Understanding by Generative Pre-Training
+# (https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf)
+class CausalLMLoss(nn.Module):
+    def __init__(self, shift_labels: bool = True) -> None:
+        super().__init__()
+        self.shift_labels = shift_labels
+        self.loss_fct = nn.CrossEntropyLoss()
+
+    def forward(
+        self, logits: torch.FloatTensor, labels: torch.LongTensor
+    ) -> torch.FloatTensor:
+        if self.shift_labels:
+            logits, labels = logits[..., :-1, :], labels[..., 1:]
+        return self.loss_fct(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))
+
+
+class PhiPreTrainedModel(PreTrainedModel):
+    config_class = PhiConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["ParallelBlock"]
+
+    def __init__(self, *inputs, **kwargs) -> None:
+        super().__init__(*inputs, **kwargs)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor = None,
+        inputs_embeds: torch.FloatTensor = None,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError(
+                "You have to specify either `input_ids` or `inputs_embeds`."
+            )
+
+        max_batch_size = (
+            inputs_embeds.shape[0] if inputs_embeds is not None else input_ids.shape[0]
+        )
+        seqlen_offset = (
+            inputs_embeds.shape[1] + input_ids.shape[1] - 2
+            if inputs_embeds is not None
+            else input_ids.shape[1] - 1
+        )
+
+        args = (
+            {"inputs_embeds": inputs_embeds}
+            if inputs_embeds is not None
+            else {"input_ids": input_ids}
+        )
+
+        if not isinstance(past_key_values, InferenceParams):
+            past_key_values = InferenceParams(
+                max_seqlen=self.config.n_positions,
+                max_batch_size=max_batch_size,
+                seqlen_offset=0,
+                batch_size_offset=0,
+                key_value_memory_dict={},
+                lengths_per_sample=None,
+            )
+        else:
+            past_key_values.seqlen_offset = seqlen_offset
+            args = {"input_ids": input_ids[:, -1].unsqueeze(-1)}
+
+        return {
+            **args,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+        }
+
+
+class PhiModel(PhiPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [""]
+    _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
+
+    def __init__(self, config: PhiConfig) -> None:
+        super().__init__(config)
+        self.embd = Embedding(config)
+        self.h = nn.ModuleList(
+            [ParallelBlock(config, block_idx=i) for i in range(config.n_layer)]
+        )
+        self.gradient_checkpointing = config.gradient_checkpointing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embd.wte
+
+    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+        self.embd.wte = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        inputs_embeds: torch.FloatTensor = None,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+    ) -> torch.FloatTensor:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of `input_ids` or `inputs_embeds`.")
+        hidden_states = self.embd(input_ids) if input_ids is not None else inputs_embeds
+
+        for layer in self.h:
+            func = layer.__call__ if self.gradient_checkpointing else layer
+            args = (hidden_states, past_key_values, attention_mask)
+            hidden_states = (
+                torch.utils.checkpoint.checkpoint(func, *args, use_reentrant=True)
+                if self.gradient_checkpointing
+                else func(*args)
+            )
+
+        return hidden_states
+
+
+class PhiForCausalLM(PhiPreTrainedModel):
+    _keys_to_ignore_on_load_missing, _keys_to_ignore_on_load_unexpected = (
+        [""],
+        [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"],
+    )
+
+    def __init__(self, config: PhiConfig) -> None:
+        super().__init__(config)
+        self.transformer = PhiModel(config)
+        self.lm_head = CausalLMHead(config)
+        self.loss = CausalLMLoss()
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head.linear
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head.linear = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        inputs_embeds: torch.FloatTensor = None,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states = self.transformer(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+        )
+        lm_logits = self.lm_head(hidden_states)
+        loss = self.loss(lm_logits, labels) if labels is not None else None
+
+        return CausalLMOutputWithPast(
+            loss=loss, logits=lm_logits, past_key_values=past_key_values
+        )

moondream/moondream.py

@@ -0,0 +1,107 @@
+import torch
+from torch import nn
+from .vision_encoder import VisionEncoder
+from .configuration_moondream import MoondreamConfig
+from transformers import PreTrainedModel
+import re
+
+from .modeling_phi import PhiForCausalLM
+from .configuration_moondream import PhiConfig
+
+class Moondream(PreTrainedModel):
+    config_class = MoondreamConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.vision_encoder = VisionEncoder()
+
+        if type(config.phi_config) == dict:
+            phi_config = PhiConfig(**config.phi_config)
+        else:
+            phi_config = config.phi_config
+        self.text_model = PhiForCausalLM(phi_config)
+
+    @property
+    def device(self):
+        return self.text_model.device
+
+    def encode_image(self, image):
+        return self.vision_encoder(image)
+
+    def input_embeds(self, prompt, image_embeds, tokenizer):
+        def _tokenize(txt):
+            return tokenizer(
+                txt, return_tensors="pt", add_special_tokens=False
+            ).input_ids.to(self.device)
+
+        text_emb = self.text_model.get_input_embeddings()
+
+        # Add BOS token
+        embeds = []
+        embeds.append(
+            text_emb((torch.tensor([[tokenizer.bos_token_id]], device=self.device)))
+        )
+
+        if "<image>" not in prompt:
+            embeds.append(text_emb(_tokenize(prompt)))
+        else:
+            assert prompt.count("<image>") == 1
+            before, after = prompt.split("<image>")
+            embeds.append(text_emb(_tokenize(f"{before}<image>")))
+            embeds.append(image_embeds.to(self.device))
+            embeds.append(text_emb(_tokenize(f"</image>{after}")))
+
+        return torch.cat(embeds, dim=1)
+
+    def generate(
+        self,
+        image_embeds,
+        prompt,
+        tokenizer,
+        eos_text="<END>",
+        max_new_tokens=128,
+        **kwargs,
+    ):
+        eos_tokens = tokenizer(eos_text, add_special_tokens=False)[0].ids
+
+        generate_config = {
+            "eos_token_id": eos_tokens,
+            "bos_token_id": tokenizer.bos_token_id,
+            "pad_token_id": tokenizer.eos_token_id,
+            "max_new_tokens": max_new_tokens,
+            **kwargs,
+        }
+
+        with torch.no_grad():
+            inputs_embeds = self.input_embeds(prompt, image_embeds, tokenizer)
+            output_ids = self.text_model.generate(
+                inputs_embeds=inputs_embeds, **generate_config
+            )
+
+        return tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+
+    def answer_question(
+        self,
+        image_embeds,
+        question,
+        tokenizer,
+        chat_history="",
+        result_queue=None,
+        **kwargs,
+    ):
+        prompt = f"<image>\n\n{chat_history}Question: {question}\n\nAnswer: "
+        answer = self.generate(
+            image_embeds,
+            prompt,
+            eos_text="<END>",
+            tokenizer=tokenizer,
+            max_new_tokens=256,
+            **kwargs,
+        )[0]
+        cleaned_answer = re.sub("<$", "", re.sub("END$", "", answer)).strip()
+
+        # Use the result_queue to pass the result if it is provided
+        if result_queue:
+            result_queue.put(cleaned_answer)
+        else:
+            return cleaned_answer

moondream/util.py

@@ -0,0 +1,13 @@
+import torch
+
+
+def detect_device():
+    """
+    Detects the appropriate device to run on, and return the device and dtype.
+    """
+    if torch.cuda.is_available():
+        return torch.device("cuda"), torch.float16
+    elif torch.backends.mps.is_available():
+        return torch.device("mps"), torch.float16
+    else:
+        return torch.device("cpu"), torch.float32

moondream/vision_encoder.py

@@ -0,0 +1,136 @@
+import torch
+from torch import nn
+from PIL import Image
+from einops import rearrange
+from torchvision.transforms.v2 import (
+    Compose,
+    Resize,
+    InterpolationMode,
+    ToImage,
+    ToDtype,
+    Normalize,
+)
+import timm
+
+
+class VisualHolder(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.visual = model
+
+    def forward(self, x):
+        return self.visual(x)
+
+
+class ModelHolder(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class LinearPatchEmbedding(nn.Module):
+    def __init__(self, conv):
+        super().__init__()
+        self.linear = nn.Linear(588, 1152)
+        self.linear.weight.data = conv.weight.data.view(1152, -1)
+        if conv.bias is not None:
+            self.linear.bias.data = conv.bias.data
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int = None,
+        out_features: int = None,
+        act_layer: nn.Module = nn.GELU,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+
+        torch.nn.init.kaiming_normal_(
+            self.fc1.weight, mode="fan_in", nonlinearity="relu"
+        )
+        torch.nn.init.kaiming_normal_(
+            self.fc2.weight, mode="fan_in", nonlinearity="relu"
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+
+class VisionProjection(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        image_embedding_dim = 1152
+        model_dim = 2048
+        hidden_dim = model_dim * 4
+
+        self.mlp = MLP(image_embedding_dim, hidden_dim, model_dim)
+
+    @property
+    def device(self):
+        return self.mlp.fc1.weight.device
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+class VisionEncoder(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.encoder = ModelHolder(
+            VisualHolder(timm.create_model("vit_so400m_patch14_siglip_384"))
+        )
+        self.encoder.model.visual.patch_embed = LinearPatchEmbedding(
+            self.encoder.model.visual.patch_embed.proj
+        )
+        self.encoder.model.visual.attn_pool = nn.Identity()
+
+        self.projection = VisionProjection()
+
+        self.preprocess = Compose(
+            [
+                Resize(size=(378, 378), interpolation=InterpolationMode.BICUBIC),
+                ToImage(),
+                ToDtype(torch.float32, scale=True),
+                Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+            ]
+        )
+
+    @property
+    def device(self):
+        return self.projection.mlp.fc1.weight.device
+
+    @property
+    def dtype(self):
+        return self.projection.mlp.fc1.weight.dtype
+
+    def __call__(self, image: Image) -> torch.Tensor:
+        with torch.no_grad():
+            x = (
+                self.preprocess(image.convert("RGB"))
+                .unsqueeze(0)
+                .to(self.device, dtype=self.dtype)
+            )
+            x = rearrange(x, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14)
+
+            x = self.encoder(x)
+            x = self.projection(x)
+
+            return x

omnichat.py

@@ -0,0 +1,219 @@
+import os
+import torch
+import json
+from PIL import Image
+import base64
+import io
+from accelerate import load_checkpoint_and_dispatch, init_empty_weights
+from transformers import AutoTokenizer, AutoModel
+
+from omnilmm.utils import disable_torch_init
+from omnilmm.model.omnilmm import OmniLMMForCausalLM
+from omnilmm.model.utils import build_transform
+from omnilmm.train.train_utils import omni_preprocess
+
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+
+    
+
+def init_omni_lmm(model_path):
+    torch.backends.cuda.matmul.allow_tf32 = True
+    disable_torch_init()
+    model_name = os.path.expanduser(model_path)
+    print(f'Load omni_lmm model and tokenizer from {model_name}')
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name, model_max_length=2048)
+
+    if False:
+        # model on multiple devices for small size gpu memory (Nvidia 3090 24G x2) 
+        with init_empty_weights():
+            model = OmniLMMForCausalLM.from_pretrained(model_name, tune_clip=True, torch_dtype=torch.bfloat16)
+        model = load_checkpoint_and_dispatch(model, model_name, dtype=torch.bfloat16, 
+                    device_map="auto",  no_split_module_classes=['Eva','MistralDecoderLayer', 'ModuleList', 'Resampler']
+        )
+    else:
+        model = OmniLMMForCausalLM.from_pretrained(
+            model_name, tune_clip=True, torch_dtype=torch.bfloat16
+        ).to(device='cuda', dtype=torch.bfloat16)
+
+    image_processor = build_transform(
+        is_train=False, input_size=model.model.config.image_size, std_mode='OPENAI_CLIP')
+
+    mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+    assert mm_use_im_start_end
+
+    tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN,
+                         DEFAULT_IM_END_TOKEN], special_tokens=True)
+
+
+    vision_config = model.model.vision_config
+    vision_config.im_patch_token = tokenizer.convert_tokens_to_ids(
+        [DEFAULT_IMAGE_PATCH_TOKEN])[0]
+    vision_config.use_im_start_end = mm_use_im_start_end
+    vision_config.im_start_token, vision_config.im_end_token = tokenizer.convert_tokens_to_ids(
+        [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN])
+    image_token_len = model.model.config.num_query
+
+    return model, image_processor, image_token_len, tokenizer
+
+def expand_question_into_multimodal(question_text, image_token_len, im_st_token, im_ed_token, im_patch_token):
+    if '<image>' in question_text[0]['content']:
+        question_text[0]['content'] = question_text[0]['content'].replace(
+            '<image>', im_st_token + im_patch_token * image_token_len + im_ed_token)
+    else:
+        question_text[0]['content'] = im_st_token + im_patch_token * \
+            image_token_len + im_ed_token + '\n' + question_text[0]['content']
+    return question_text
+
+def wrap_question_for_omni_lmm(question, image_token_len, tokenizer):
+    question = expand_question_into_multimodal(
+        question, image_token_len, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN)
+
+    conversation = question
+    data_dict = omni_preprocess(sources=[conversation],
+                                  tokenizer=tokenizer,
+                                  generation=True)
+
+    data_dict = dict(input_ids=data_dict["input_ids"][0],
+                     labels=data_dict["labels"][0])
+    return data_dict
+
+
+
+class OmniLMM12B:
+    def __init__(self, model_path) -> None:
+        model, img_processor, image_token_len, tokenizer = init_omni_lmm(model_path)
+        self.model = model
+        self.image_token_len = image_token_len
+        self.image_transform = img_processor
+        self.tokenizer = tokenizer
+        self.model.eval()
+
+    def decode(self, image, input_ids):
+        with torch.inference_mode():
+            output = self.model.generate_vllm(
+                input_ids=input_ids.unsqueeze(0).cuda(),
+                images=image.unsqueeze(0).half().cuda(),
+                temperature=0.6,
+                max_new_tokens=1024,
+                # num_beams=num_beams,
+                do_sample=True,
+                output_scores=True,
+                return_dict_in_generate=True,
+                repetition_penalty=1.1,
+                top_k=30,
+                top_p=0.9,
+            )
+
+            response = self.tokenizer.decode(
+                output.sequences[0], skip_special_tokens=True)
+            response = response.strip()
+            return response
+
+    def chat(self, input):
+        try:
+            image = Image.open(io.BytesIO(base64.b64decode(input['image']))).convert('RGB')
+        except Exception as e:
+            return "Image decode error"
+
+        msgs = json.loads(input['question'])
+        input_ids = wrap_question_for_omni_lmm(
+            msgs, self.image_token_len, self.tokenizer)['input_ids']
+        input_ids = torch.as_tensor(input_ids)
+        #print('input_ids', input_ids)
+        image = self.image_transform(image)
+
+        out = self.decode(image, input_ids)
+
+        return out
+        
+
+def img2base64(file_name):
+    with open(file_name, 'rb') as f:
+        encoded_string = base64.b64encode(f.read())
+        return encoded_string
+
+class OmniLMM3B:
+    def __init__(self, model_path) -> None:
+        self.model = AutoModel.from_pretrained(model_path, trust_remote_code=True).to(dtype=torch.bfloat16)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model.eval().cuda()
+
+    def chat(self, input):
+        try:
+            image = Image.open(io.BytesIO(base64.b64decode(input['image']))).convert('RGB')
+        except Exception as e:
+            return "Image decode error"
+
+        msgs = json.loads(input['question'])
+        
+        answer, context, _ = self.model.chat(
+            image=image,
+            msgs=msgs,
+            context=None,
+            tokenizer=self.tokenizer,
+            sampling=True,
+            temperature=0.7
+    	)
+        return answer
+
+class MiniCPMV2_5:
+    def __init__(self, model_path) -> None:
+        self.model = AutoModel.from_pretrained(model_path, trust_remote_code=True).to(dtype=torch.float16)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model.eval().cuda()
+
+    def chat(self, input):
+        try:
+            image = Image.open(io.BytesIO(base64.b64decode(input['image']))).convert('RGB')
+        except Exception as e:
+            return "Image decode error"
+
+        msgs = json.loads(input['question'])
+        
+        answer = self.model.chat(
+            image=image,
+            msgs=msgs,
+            tokenizer=self.tokenizer,
+            sampling=True,
+            temperature=0.7
+    	)
+        return answer
+
+
+class OmniLMMChat:
+    def __init__(self, model_path) -> None:
+        if '12B' in model_path:
+            self.model = OmniLMM12B(model_path)
+        elif 'MiniCPM-Llama3-V' in model_path:
+            self.model = MiniCPMV2_5(model_path)
+        else:
+            self.model = OmniLMM3B(model_path)
+
+    def chat(self, input):
+        return self.model.chat(input)
+
+
+if __name__ == '__main__':
+    
+    model_path = 'openbmb/OmniLMM-12B'
+    chat_model = OmniLMMChat(model_path)
+
+    im_64 = img2base64('./assets/worldmap_ck.jpg')
+
+    # first round chat 
+    msgs = [{"role": "user", "content": "What is interesting about this image?"}]
+    input = {"image": im_64, "question": json.dumps(msgs, ensure_ascii=True)}
+    answer = chat_model.chat(input)
+    print(msgs[-1]["content"]+'\n', answer)
+
+    # second round chat 
+    msgs.append({"role": "assistant", "content": answer})
+    msgs.append({"role": "user", "content": "Where is China in the image"})
+    input = {"image": im_64,"question": json.dumps(msgs, ensure_ascii=True)}
+    answer = chat_model.chat(input)
+    print(msgs[-1]["content"]+'\n', answer)
+

omnilmm/constants.py

@@ -0,0 +1,4 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."

omnilmm/conversation.py

@@ -0,0 +1,320 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(
+                                    pil_img.mode, (width, width), background_color)
+                                result.paste(
+                                    pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(
+                                    pil_img.mode, (height, height), background_color)
+                                result.paste(
+                                    pil_img, ((height - width) // 2, 0))
+                                return result
+                        image = expand2square(image)
+                    elif image_process_mode == "Crop":
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((224, 224))
+                    else:
+                        raise ValueError(
+                            f"Invalid image_process_mode: {image_process_mode}")
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(
+                        min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="JPEG")
+                        img_b64_str = base64.b64encode(
+                            buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    msg, image, image_process_mode = msg
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(
+                        min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    # image = image.resize((224, 224))
+                    buffered = BytesIO()
+                    image.save(buffered, format="JPEG")
+                    img_b64_str = base64.b64encode(
+                        buffered.getvalue()).decode()
+                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = msg.replace('<image>', img_str)
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "Give three tips for staying healthy."),
+        ("Assistant",
+            "Sure, here are three tips for staying healthy:\n"
+            "1. Exercise regularly: Regular physical activity can help improve your overall health and wellbeing. "
+            "It can also help reduce your risk of chronic conditions such as obesity, diabetes, heart disease, "
+            "and certain cancers. Aim for at least 150 minutes of moderate-intensity aerobic exercise or "
+            "75 minutes of vigorous-intensity aerobic exercise per week, along with muscle-strengthening "
+            "activities at least two days per week.\n"
+            "2. Eat a balanced diet: Eating a balanced diet that is rich in fruits, "
+            "vegetables, whole grains, lean proteins, and healthy fats can help support "
+            "your overall health. Try to limit your intake of processed and high-sugar foods, "
+            "and aim to drink plenty of water throughout the day.\n"
+            "3. Get enough sleep: Getting enough quality sleep is essential for your physical "
+            "and mental health. Adults should aim for seven to nine hours of sleep per night. "
+            "Establish a regular sleep schedule and try to create a relaxing bedtime routine to "
+            "help improve the quality of your sleep.")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_v1_2 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1_1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_bair_v1 = Conversation(
+    system="BEGINNING OF CONVERSATION:",
+    roles=("USER", "GPT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+simple_conv = Conversation(
+    system="You are LLaVA, a large language model trained by UW Madison WAIV Lab, based on LLaMA architecture."
+           "You are designed to assist human with a variety of tasks using natural language."
+           "Follow the instructions carefully.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "Hi!"),
+        ("Assistant", "Hi there!  How can I help you today?\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+simple_conv_multimodal = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+simple_conv_legacy = Conversation(
+    system="You are LLaVA, a large language model trained by UW Madison WAIV Lab."
+           "You are designed to assist human with a variety of tasks using natural language."
+           "Follow the instructions carefully.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "Hi!\n\n### Response:"),
+        ("Assistant", "Hi there!  How can I help you today?\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v1 = Conversation(
+    system="You are LLaVA, a large language and vision assistant trained by UW Madison WAIV Lab."
+           "You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "Follow the instructions carefully and explain your answers in detail.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+default_conversation = conv_v1_2
+conv_templates = {
+    "default": conv_v1_2,
+    "simple": simple_conv,
+    "simple_legacy": simple_conv_legacy,
+    "multimodal": simple_conv_multimodal,
+    "llava_v1": conv_llava_v1,
+
+    # fastchat
+    "v1": conv_v1_2,
+    "bair_v1": conv_bair_v1,
+    "vicuna_v1_1": conv_vicuna_v1_1,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

omnilmm/model/__init__.py

@@ -0,0 +1 @@
+from .omnilmm import OmniLMMForCausalLM

omnilmm/model/omnilmm.py

@@ -0,0 +1,457 @@
+
+import gc
+import math
+import timm
+import torch
+from torch import Tensor
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+from typing import List, Optional, Tuple, Union
+
+from transformers import AutoConfig, AutoModelForCausalLM
+from transformers import MistralForCausalLM, MistralModel, MistralConfig
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from omnilmm.model.utils import build_transform
+from omnilmm.model.resampler import Resampler
+
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+
+
+class OmniLMMConfig(MistralConfig):
+    model_type = "omnilmm"
+
+
+class Identity(torch.nn.Identity):
+    def forward(self, input: Tensor, **kwargs) -> Tensor:
+        return super().forward(input)
+
+
+def create_vision_module(config):
+    vision_tower = timm.create_model('eva02_enormous_patch14_clip_224.laion2b_plus',
+                                     pretrained=False,
+                                     num_classes=0,
+                                     dynamic_img_size=True,
+                                     dynamic_img_pad=True)
+
+    if isinstance(vision_tower, timm.models.VisionTransformer):
+        if vision_tower.attn_pool is not None:
+            vision_tower.attn_pool = Identity()
+
+    # use 2nd last layer's output
+    vision_tower.blocks[-1] = Identity()
+
+    embed_dim = config.hidden_size
+    resampler = Resampler(
+        grid_size=int(math.sqrt(config.num_query)),
+        embed_dim=embed_dim,
+        num_heads=embed_dim // 128,
+        kv_dim=vision_tower.embed_dim,
+    )
+    return vision_tower, resampler
+
+
+class OmniLMMModel(MistralModel):
+    config_class = OmniLMMConfig
+
+    def __init__(self, config: OmniLMMConfig, mm_vision_tower=None, mm_hidden_size=None, tune_clip=True):
+        super(OmniLMMModel, self).__init__(config)
+
+        if hasattr(config, "mm_vision_tower"):
+            vision_tower, resampler = create_vision_module(config)
+
+            # print(__file__, 'skip loading vision tower weights')
+
+            # HACK: for FSDP
+            self.vision_tower = [vision_tower]
+            self.resampler = resampler
+            if tune_clip:
+                self.vision_tower = self.vision_tower[0]
+
+        self.vision_config = lambda x: None
+
+    def initialize_vision_modules(self, vision_tower, no_randaug, num_query, image_size, tune_clip=False):
+        self.config.mm_vision_tower = vision_tower
+        self.config.use_mm_proj = True
+        self.config.num_query = num_query
+        self.config.image_size = image_size
+
+        if not hasattr(self, 'vision_tower'):
+            vision_tower, resampler = create_vision_module(self.config)
+            state_dict = torch.load(
+                '/tt/data/public/multimodal/multimodal_model_ckpts/timm/eva02_enormous_patch14_clip_224.laion2b_plus.pt')
+            vision_tower.load_state_dict(state_dict, strict=False)
+            del state_dict
+            gc.collect()
+        else:
+            if isinstance(self.vision_tower, list):
+                vision_tower = self.vision_tower[0]
+            else:
+                vision_tower = self.vision_tower
+            resampler = self.resampler
+        self.vision_tower = vision_tower if tune_clip else [vision_tower]
+        self.resampler = resampler
+
+        train_img_transform = build_transform(
+            is_train=True, randaug=not no_randaug, input_size=self.config.image_size, std_mode='OPENAI_CLIP')
+        eval_img_transform = build_transform(
+            is_train=False, input_size=self.config.image_size, std_mode='OPENAI_CLIP')
+
+        return dict(
+            image_processor=(train_img_transform, eval_img_transform),
+            image_token_len=num_query,
+            vision_config=self.vision_config
+        )
+
+    def get_vision_embedding(self, pixel_values):
+        if isinstance(self.vision_tower, list):
+            vision_tower = self.vision_tower[0]  # HACK: for FSDP
+        else:
+            vision_tower = self.vision_tower
+
+        dtype = vision_tower.pos_embed.data.dtype
+        vision_embedding = vision_tower.forward_features(
+            pixel_values.type(dtype))
+        if hasattr(vision_tower, 'num_prefix_tokens') and vision_tower.num_prefix_tokens > 0:
+            vision_embedding = vision_embedding[:,
+                                                vision_tower.num_prefix_tokens:]
+        res = self.resampler(vision_embedding)
+        return res
+
+    def get_vllm_embedding(self, data):
+
+        if 'vision_hidden_states' not in data:
+            pixel_values_list = data['pixel_values']
+            vision_hidden_states = []
+            for pixel_values in pixel_values_list:
+                if len(pixel_values) > 0:
+                    vision_hidden_states.append(self.get_vision_embedding(pixel_values.unsqueeze(0))[0])
+                else:
+                    vision_hidden_states.append([])
+        else:
+            vision_hidden_states = data['vision_hidden_states']
+
+        #vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
+        inputs_embeds = self.embed_tokens(data['input_ids'])
+        vision_hidden_states = [i.type(inputs_embeds.dtype) 
+            if isinstance(i, torch.Tensor) else i for i in vision_hidden_states
+        ]
+
+
+        # HACK: replace back original embeddings for LLaVA pretraining
+        orig_embeds_params = getattr(self, 'orig_embeds_params', None)
+
+        new_input_embeds = []
+        cur_image_idx = 0
+        for cur_input_ids, cur_input_embeds in zip(data['input_ids'], inputs_embeds):
+            if (cur_input_ids == self.vision_config.im_patch_token).sum() == 0:
+                # multimodal LLM, but the current sample is not multimodal
+                cur_input_embeds = cur_input_embeds + (0. * dummy_image_features).sum()
+                new_input_embeds.append(cur_input_embeds)
+                continue
+
+            if self.vision_config.use_im_start_end:
+                cur_image_features = vision_hidden_states[cur_image_idx]
+                num_patches = cur_image_features.shape[0]
+                if (cur_input_ids == self.vision_config.im_start_token).sum() != (cur_input_ids == self.vision_config.im_end_token).sum():
+                    raise ValueError(
+                        "The number of image start tokens and image end tokens should be the same.")
+                image_start_tokens = torch.where(
+                    cur_input_ids == self.vision_config.im_start_token)[0]
+                for image_start_token_pos in image_start_tokens:
+                    cur_image_features = vision_hidden_states[cur_image_idx].to(
+                        device=cur_input_embeds.device)
+                    num_patches = cur_image_features.shape[0]
+                    if cur_input_ids[image_start_token_pos + num_patches + 1] != self.vision_config.im_end_token:
+                        raise ValueError(
+                            "The image end token should follow the image start token.")
+                    if orig_embeds_params is not None:
+                        cur_new_input_embeds = torch.cat((cur_input_embeds[:image_start_token_pos].detach(), cur_input_embeds[image_start_token_pos:image_start_token_pos+1], cur_image_features,
+                                                         cur_input_embeds[image_start_token_pos + num_patches + 1:image_start_token_pos + num_patches + 2], cur_input_embeds[image_start_token_pos + num_patches + 2:].detach()), dim=0)
+                    else:
+                        cur_new_input_embeds = torch.cat(
+                            (cur_input_embeds[:image_start_token_pos+1], cur_image_features, cur_input_embeds[image_start_token_pos + num_patches + 1:]), dim=0)
+                    cur_image_idx += 1
+                new_input_embeds.append(cur_new_input_embeds)
+            else:
+                raise NotImplementedError
+        inputs_embeds = torch.stack(new_input_embeds, dim=0)
+
+        return inputs_embeds, vision_hidden_states
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+
+        # HACK: replace back original embeddings for LLaVA pretraining
+        orig_embeds_params = getattr(self, 'orig_embeds_params', None)
+
+        if inputs_embeds is None and past_key_values is None:
+          inputs_embeds = self.embed_tokens(input_ids)
+
+          vision_tower = getattr(self, 'vision_tower', None)
+          if vision_tower is not None and (input_ids.shape[1] != 1 or self.training) and images is not None:
+
+            if type(images) is list:
+                image_features = []
+                for image in images:
+                    image_forward_out = self.get_vision_embedding(image.unsqueeze(0))[
+                        0]
+                    image_features.append(image_forward_out)
+            else:
+                image_features = self.get_vision_embedding(images)
+
+            dummy_image_features = torch.zeros(
+                self.config.num_query,
+                self.config.hidden_size,
+                device=inputs_embeds.device,
+                dtype=inputs_embeds.dtype)
+
+            new_input_embeds = []
+            cur_image_idx = 0
+            for cur_input_ids, cur_input_embeds in zip(input_ids, inputs_embeds):
+                if (cur_input_ids == self.vision_config.im_patch_token).sum() == 0:
+                    # multimodal LLM, but the current sample is not multimodal
+                    cur_input_embeds = cur_input_embeds + \
+                        (0. * dummy_image_features).sum()
+                    new_input_embeds.append(cur_input_embeds)
+                    continue
+
+                if self.vision_config.use_im_start_end:
+                    cur_image_features = image_features[cur_image_idx]
+                    num_patches = cur_image_features.shape[0]
+                    if (cur_input_ids == self.vision_config.im_start_token).sum() != (cur_input_ids == self.vision_config.im_end_token).sum():
+                        raise ValueError(
+                            "The number of image start tokens and image end tokens should be the same.")
+                    image_start_tokens = torch.where(
+                        cur_input_ids == self.vision_config.im_start_token)[0]
+                    for image_start_token_pos in image_start_tokens:
+                        cur_image_features = image_features[cur_image_idx].to(
+                            device=cur_input_embeds.device)
+                        num_patches = cur_image_features.shape[0]
+                        if cur_input_ids[image_start_token_pos + num_patches + 1] != self.vision_config.im_end_token:
+                            raise ValueError(
+                                "The image end token should follow the image start token.")
+                        if orig_embeds_params is not None:
+                            cur_new_input_embeds = torch.cat((cur_input_embeds[:image_start_token_pos].detach(), cur_input_embeds[image_start_token_pos:image_start_token_pos+1], cur_image_features,
+                                                             cur_input_embeds[image_start_token_pos + num_patches + 1:image_start_token_pos + num_patches + 2], cur_input_embeds[image_start_token_pos + num_patches + 2:].detach()), dim=0)
+                        else:
+                            cur_new_input_embeds = torch.cat(
+                                (cur_input_embeds[:image_start_token_pos+1], cur_image_features, cur_input_embeds[image_start_token_pos + num_patches + 1:]), dim=0)
+                        cur_image_idx += 1
+                    new_input_embeds.append(cur_new_input_embeds)
+                else:
+                    raise NotImplementedError
+            inputs_embeds = torch.stack(new_input_embeds, dim=0)
+            input_ids = None
+
+        return super(OmniLMMModel, self).forward(
+            input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds, use_cache=use_cache,
+            output_attentions=output_attentions, output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs
+        )
+
+
+class OmniLMMForCausalLM(MistralForCausalLM):
+    config_class = OmniLMMConfig
+
+    def __init__(self, config, mm_vision_tower=None, tune_clip=True):
+        super(MistralForCausalLM, self).__init__(config)
+        self.model = OmniLMMModel(
+            config, mm_vision_tower=mm_vision_tower, tune_clip=tune_clip)
+
+        self.lm_head = nn.Linear(
+            config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # print(f'@@@ At forward, labels: {labels.shape}-{labels}', flush=True)
+        # print(f'@@@ At forward, input_ids: {input_ids.shape}-{input_ids}', flush=True)
+        # print(f'@@@ At forward, input_ids: {attention_mask.shape}-{attention_mask}', flush=True)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            images=images,
+            **kwargs
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model/pipeline parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    # TODO could be removed for generate_vllm()
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "images": kwargs.get("images", None),
+            }
+        )
+        return model_inputs
+
+    def generate_vllm(
+        self,
+        input_ids: torch.LongTensor = None,
+        images: Optional[torch.FloatTensor] = None,
+        vision_hidden_states=None,
+        return_vision_hidden_states=False,
+        **kwargs
+    ):
+        model_inputs = {'input_ids': input_ids}
+        if vision_hidden_states is None:
+            model_inputs['pixel_values'] = images
+        else:
+            model_inputs['vision_hidden_states'] = vision_hidden_states
+
+        with torch.inference_mode():
+            inputs_embeds, vision_hidden_states = self.model.get_vllm_embedding(model_inputs)
+
+            result = self.generate(
+                inputs_embeds=inputs_embeds,
+                **kwargs
+            )
+
+        if return_vision_hidden_states:
+            return result, vision_hidden_states
+
+        return result
+
+
+    def initialize_vision_tokenizer(self, mm_use_im_start_end, tokenizer, device,
+                                    tune_mm_mlp_adapter=False):
+        self.model.vision_config.use_im_start_end = mm_use_im_start_end
+        tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+        self.resize_token_embeddings(len(tokenizer))
+
+        if mm_use_im_start_end:
+            num_new_tokens = tokenizer.add_tokens(
+                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+            self.model.vision_config.im_start_token, self.model.vision_config.im_end_token = tokenizer.convert_tokens_to_ids(
+                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN])
+
+            if num_new_tokens > 0:
+                input_embeddings = self.get_input_embeddings().weight.data
+                output_embeddings = self.get_output_embeddings().weight.data
+
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+
+                input_embeddings[-num_new_tokens:] = input_embeddings_avg
+                output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+            # for new sft data
+            num_new_tokens = tokenizer.add_tokens(
+                ['<box>', '</box>', '<ref>', '</ref>', '<quad>', '</quad>'], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+            if num_new_tokens > 0:
+                input_embeddings = self.get_input_embeddings().weight.data
+                output_embeddings = self.get_output_embeddings().weight.data
+
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+
+                input_embeddings[-num_new_tokens:] = input_embeddings_avg
+                output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+            if tune_mm_mlp_adapter:
+                self.model.orig_embeds_params = [
+                    self.get_input_embeddings().weight.data.clone().to(device=device)]
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = True
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False
+
+        self.model.vision_config.im_patch_token = tokenizer.convert_tokens_to_ids(
+            [DEFAULT_IMAGE_PATCH_TOKEN])[0]
+        print(f'Tokenizer: {tokenizer}\n patch_token_id: {self.model.vision_config.im_patch_token}, visoin_config: {self.model.vision_config}', flush=True)
+        # exit()
+
+
+AutoConfig.register("omnilmm", OmniLMMConfig)
+AutoModelForCausalLM.register(OmniLMMConfig, OmniLMMForCausalLM)

omnilmm/model/resampler.py

@@ -0,0 +1,171 @@
+# Copyright (c) Alibaba Cloud.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import OrderedDict
+import math
+import requests
+from io import BytesIO
+from functools import partial
+from PIL import Image
+from typing import Callable, Optional, Sequence, Tuple, List, Union
+import numpy as np
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.init import trunc_normal_
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+
+
+def get_abs_pos(abs_pos, tgt_size):
+    # abs_pos: L, C
+    # tgt_size: M
+    # return: M, C
+    src_size = int(math.sqrt(abs_pos.size(0)))
+    tgt_size = int(math.sqrt(tgt_size))
+    dtype = abs_pos.dtype
+
+    if src_size != tgt_size:
+        return F.interpolate(
+            abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2),
+            size=(tgt_size, tgt_size),
+            mode="bicubic",
+            align_corners=False,
+        ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype)
+    else:
+        return abs_pos
+
+
+# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate(
+            [np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(
+        embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(
+        embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000 ** omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+class Resampler(nn.Module):
+    """
+    A 2D perceiver-resampler network with one cross attention layers by
+        (grid_size**2) learnable queries and 2d sincos pos_emb
+    Outputs:
+        A tensor with the shape of (grid_size**2, embed_dim)
+    """
+
+    def __init__(
+            self,
+            grid_size,
+            embed_dim,
+            num_heads,
+            kv_dim=None,
+            norm_layer=partial(nn.LayerNorm, eps=1e-6)
+    ):
+        super().__init__()
+        self.num_queries = grid_size ** 2
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+
+        self.pos_embed = nn.Parameter(
+            torch.from_numpy(get_2d_sincos_pos_embed(
+                embed_dim, grid_size)).float()
+        ).requires_grad_(False)
+
+        self.query = nn.Parameter(torch.zeros(self.num_queries, embed_dim))
+        trunc_normal_(self.query, std=.02)
+
+        if kv_dim is not None and kv_dim != embed_dim:
+            self.kv_proj = nn.Linear(kv_dim, embed_dim, bias=False)
+        else:
+            self.kv_proj = nn.Identity()
+
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
+        self.ln_q = norm_layer(embed_dim)
+        self.ln_kv = norm_layer(embed_dim)
+
+        self.ln_post = norm_layer(embed_dim)
+        self.proj = nn.Parameter(
+            (embed_dim ** -0.5) * torch.randn(embed_dim, embed_dim))
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x, attn_mask=None):
+
+        pos_embed = get_abs_pos(self.pos_embed, x.size(1))
+
+        x = self.kv_proj(x)
+        x = self.ln_kv(x).permute(1, 0, 2)
+
+        N = x.shape[1]
+        q = self.ln_q(self.query)
+        # print((self._repeat(q, N) + self.pos_embed.unsqueeze(1)).dtype, (x + pos_embed.unsqueeze(1)).dtype, x.dtype)
+        out = self.attn(
+            self._repeat(q, N) + self.pos_embed.unsqueeze(1),
+            x + pos_embed.unsqueeze(1),
+            x,
+            attn_mask=attn_mask)[0]
+        x = out.permute(1, 0, 2)
+
+        x = self.ln_post(x)
+        x = x @ self.proj
+        return x
+
+    def _repeat(self, query, N: int):
+        return query.unsqueeze(1).repeat(1, N, 1)

omnilmm/model/utils.py

@@ -0,0 +1,555 @@
+from torchvision import transforms
+from timm.data.transforms import RandomResizedCropAndInterpolation
+from timm.data.constants import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
+from transformers import AutoConfig
+from PIL import Image
+from io import BytesIO
+import torch.distributed as dist
+import numpy as np
+import pickle
+import base64
+import cv2
+import os
+import torch
+from transformers import AutoConfig, StoppingCriteria
+
+try:
+    from timm.data.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
+except ImportError:
+    OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
+    OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
+
+
+def auto_upgrade(config):
+    cfg = AutoConfig.from_pretrained(config)
+    if 'llava' in config and cfg.model_type != 'llava':
+        print("You are using newer LLaVA code base, while the checkpoint of v0 is from older code base.")
+        print("You must upgrade the checkpoint to the new code base (this can be done automatically).")
+        confirm = input(
+            "Please confirm that you want to upgrade the checkpoint. [Y/N]")
+        if confirm.lower() in ["y", "yes"]:
+            print("Upgrading checkpoint...")
+            assert len(cfg.architectures) == 1
+            setattr(cfg.__class__, "model_type", "llava")
+            cfg.architectures[0] = 'LlavaLlamaForCausalLM'
+            cfg.save_pretrained(config)
+            print("Checkpoint upgraded.")
+        else:
+            print("Checkpoint upgrade aborted.")
+            exit(1)
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.tokenizer = tokenizer
+        self.start_len = None
+        self.input_ids = input_ids
+
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        if self.start_len is None:
+            self.start_len = self.input_ids.shape[1]
+        else:
+            outputs = self.tokenizer.batch_decode(
+                output_ids[:, self.start_len:], skip_special_tokens=True)[0]
+            for keyword in self.keywords:
+                if keyword in outputs:
+                    return True
+        return False
+
+
+def auto_upgrade(config):
+    cfg = AutoConfig.from_pretrained(config)
+    if 'llava' in config and cfg.model_type != 'llava':
+        print("You are using newer LLaVA code base, while the checkpoint of v0 is from older code base.")
+        print("You must upgrade the checkpoint to the new code base (this can be done automatically).")
+        confirm = input(
+            "Please confirm that you want to upgrade the checkpoint. [Y/N]")
+        if confirm.lower() in ["y", "yes"]:
+            print("Upgrading checkpoint...")
+            assert len(cfg.architectures) == 1
+            setattr(cfg.__class__, "model_type", "llava")
+            cfg.architectures[0] = 'LlavaLlamaForCausalLM'
+            cfg.save_pretrained(config)
+            print("Checkpoint upgraded.")
+        else:
+            print("Checkpoint upgrade aborted.")
+            exit(1)
+
+# aug functions
+
+
+def identity_func(img):
+    return img
+
+
+def autocontrast_func(img, cutoff=0):
+    '''
+        same output as PIL.ImageOps.autocontrast
+    '''
+    n_bins = 256
+
+    def tune_channel(ch):
+        n = ch.size
+        cut = cutoff * n // 100
+        if cut == 0:
+            high, low = ch.max(), ch.min()
+        else:
+            hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+            low = np.argwhere(np.cumsum(hist) > cut)
+            low = 0 if low.shape[0] == 0 else low[0]
+            high = np.argwhere(np.cumsum(hist[::-1]) > cut)
+            high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]
+        if high <= low:
+            table = np.arange(n_bins)
+        else:
+            scale = (n_bins - 1) / (high - low)
+            table = np.arange(n_bins) * scale - low * scale
+            table[table < 0] = 0
+            table[table > n_bins - 1] = n_bins - 1
+        table = table.clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def equalize_func(img):
+    '''
+        same output as PIL.ImageOps.equalize
+        PIL's implementation is different from cv2.equalize
+    '''
+    n_bins = 256
+
+    def tune_channel(ch):
+        hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+        non_zero_hist = hist[hist != 0].reshape(-1)
+        step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)
+        if step == 0:
+            return ch
+        n = np.empty_like(hist)
+        n[0] = step // 2
+        n[1:] = hist[:-1]
+        table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def rotate_func(img, degree, fill=(0, 0, 0)):
+    '''
+    like PIL, rotate by degree, not radians
+    '''
+    H, W = img.shape[0], img.shape[1]
+    center = W / 2, H / 2
+    M = cv2.getRotationMatrix2D(center, degree, 1)
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill)
+    return out
+
+
+def solarize_func(img, thresh=128):
+    '''
+        same output as PIL.ImageOps.posterize
+    '''
+    table = np.array([el if el < thresh else 255 - el for el in range(256)])
+    table = table.clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def color_func(img, factor):
+    '''
+        same output as PIL.ImageEnhance.Color
+    '''
+    # implementation according to PIL definition, quite slow
+    #  degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
+    #  out = blend(degenerate, img, factor)
+    #  M = (
+    #      np.eye(3) * factor
+    #      + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)
+    #  )[np.newaxis, np.newaxis, :]
+    M = (
+        np.float32([
+            [0.886, -0.114, -0.114],
+            [-0.587, 0.413, -0.587],
+            [-0.299, -0.299, 0.701]]) * factor
+        + np.float32([[0.114], [0.587], [0.299]])
+    )
+    out = np.matmul(img, M).clip(0, 255).astype(np.uint8)
+    return out
+
+
+def contrast_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Contrast
+    """
+    mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))
+    table = np.array([(
+        el - mean) * factor + mean
+        for el in range(256)
+    ]).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def brightness_func(img, factor):
+    '''
+        same output as PIL.ImageEnhance.Contrast
+    '''
+    table = (np.arange(256, dtype=np.float32) *
+             factor).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def sharpness_func(img, factor):
+    '''
+    The differences the this result and PIL are all on the 4 boundaries, the center
+    areas are same
+    '''
+    kernel = np.ones((3, 3), dtype=np.float32)
+    kernel[1][1] = 5
+    kernel /= 13
+    degenerate = cv2.filter2D(img, -1, kernel)
+    if factor == 0.0:
+        out = degenerate
+    elif factor == 1.0:
+        out = img
+    else:
+        out = img.astype(np.float32)
+        degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]
+        out[1:-1, 1:-1, :] = degenerate + factor * \
+            (out[1:-1, 1:-1, :] - degenerate)
+        out = out.astype(np.uint8)
+    return out
+
+
+def shear_x_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, factor, 0], [0, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill,
+                         flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def translate_x_func(img, offset, fill=(0, 0, 0)):
+    '''
+        same output as PIL.Image.transform
+    '''
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, -offset], [0, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill,
+                         flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def translate_y_func(img, offset, fill=(0, 0, 0)):
+    '''
+        same output as PIL.Image.transform
+    '''
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [0, 1, -offset]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill,
+                         flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def posterize_func(img, bits):
+    '''
+        same output as PIL.ImageOps.posterize
+    '''
+    out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))
+    return out
+
+
+def shear_y_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [factor, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill,
+                         flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def cutout_func(img, pad_size, replace=(0, 0, 0)):
+    replace = np.array(replace, dtype=np.uint8)
+    H, W = img.shape[0], img.shape[1]
+    rh, rw = np.random.random(2)
+    pad_size = pad_size // 2
+    ch, cw = int(rh * H), int(rw * W)
+    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
+    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
+    out = img.copy()
+    out[x1:x2, y1:y2, :] = replace
+    return out
+
+
+# level to args
+def enhance_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        return ((level / MAX_LEVEL) * 1.8 + 0.1,)
+    return level_to_args
+
+
+def shear_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 0.3
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def translate_level_to_args(translate_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * float(translate_const)
+        if np.random.random() > 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * cutout_const)
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def solarize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 256)
+        return (level, )
+    return level_to_args
+
+
+def none_level_to_args(level):
+    return ()
+
+
+def posterize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 4)
+        return (level, )
+    return level_to_args
+
+
+def rotate_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 30
+        if np.random.random() < 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+func_dict = {
+    'Identity': identity_func,
+    'AutoContrast': autocontrast_func,
+    'Equalize': equalize_func,
+    'Rotate': rotate_func,
+    'Solarize': solarize_func,
+    'Color': color_func,
+    'Contrast': contrast_func,
+    'Brightness': brightness_func,
+    'Sharpness': sharpness_func,
+    'ShearX': shear_x_func,
+    'TranslateX': translate_x_func,
+    'TranslateY': translate_y_func,
+    'Posterize': posterize_func,
+    'ShearY': shear_y_func,
+}
+
+translate_const = 10
+MAX_LEVEL = 10
+replace_value = (128, 128, 128)
+arg_dict = {
+    'Identity': none_level_to_args,
+    'AutoContrast': none_level_to_args,
+    'Equalize': none_level_to_args,
+    'Rotate': rotate_level_to_args(MAX_LEVEL, replace_value),
+    'Solarize': solarize_level_to_args(MAX_LEVEL),
+    'Color': enhance_level_to_args(MAX_LEVEL),
+    'Contrast': enhance_level_to_args(MAX_LEVEL),
+    'Brightness': enhance_level_to_args(MAX_LEVEL),
+    'Sharpness': enhance_level_to_args(MAX_LEVEL),
+    'ShearX': shear_level_to_args(MAX_LEVEL, replace_value),
+    'TranslateX': translate_level_to_args(
+        translate_const, MAX_LEVEL, replace_value
+    ),
+    'TranslateY': translate_level_to_args(
+        translate_const, MAX_LEVEL, replace_value
+    ),
+    'Posterize': posterize_level_to_args(MAX_LEVEL),
+    'ShearY': shear_level_to_args(MAX_LEVEL, replace_value),
+}
+
+
+class RandomAugment(object):
+
+    def __init__(self, N=2, M=10, isPIL=False, augs=[]):
+        self.N = N
+        self.M = M
+        self.isPIL = isPIL
+        if augs:
+            self.augs = augs
+        else:
+            self.augs = list(arg_dict.keys())
+
+    def get_random_ops(self):
+        sampled_ops = np.random.choice(self.augs, self.N)
+        return [(op, 0.5, self.M) for op in sampled_ops]
+
+    def __call__(self, img):
+        if self.isPIL:
+            img = np.array(img)
+        ops = self.get_random_ops()
+        for name, prob, level in ops:
+            if np.random.random() > prob:
+                continue
+            args = arg_dict[name](level)
+            img = func_dict[name](img, *args)
+        return img
+
+
+def build_transform(is_train, randaug=True, input_size=224, interpolation='bicubic', std_mode='IMAGENET_INCEPTION'):
+    if std_mode == 'IMAGENET_INCEPTION':
+        mean = IMAGENET_INCEPTION_MEAN
+        std = IMAGENET_INCEPTION_STD
+    elif std_mode == 'OPENAI_CLIP':
+        mean = OPENAI_CLIP_MEAN
+        std = OPENAI_CLIP_STD
+    else:
+        raise NotImplementedError
+
+    if is_train:
+        crop_scale = float(os.environ.get('TRAIN_CROP_SCALE', 0.9999))
+        t = [
+            RandomResizedCropAndInterpolation(
+                input_size, scale=(crop_scale, 1.0), interpolation='bicubic'),
+            # transforms.RandomHorizontalFlip(),
+        ]
+        if randaug and os.environ.get('TRAIN_DO_AUG', 'False') == 'True':
+            print(f'@@@@@ Do random aug during training', flush=True)
+            t.append(
+                RandomAugment(
+                    2, 7, isPIL=True,
+                    augs=[
+                        'Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness',
+                        'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Rotate',
+                    ]))
+        else:
+            print(f'@@@@@ Skip random aug during training', flush=True)
+        t += [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean, std=std),
+        ]
+        t = transforms.Compose(t)
+    else:
+        t = transforms.Compose([
+            transforms.Resize((input_size, input_size),
+                              interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean, std=std)
+        ])
+
+    return t
+
+
+def img2b64(img_path):
+    img = Image.open(img_path)  # path to file
+    img_buffer = BytesIO()
+    img.save(img_buffer, format=img.format)
+    byte_data = img_buffer.getvalue()
+    base64_str = base64.b64encode(byte_data)  # bytes
+    base64_str = base64_str.decode("utf-8")  # str
+    return base64_str
+
+
+def str2b64(str):
+    return base64.b64encode(str.encode('utf-8')).decode('utf-8')
+
+
+def b642str(b64):
+    return base64.b64decode(b64).decode('utf-8')
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to("cuda")
+
+    # obtain Tensor size of each rank
+    local_size = torch.LongTensor([tensor.numel()]).to("cuda")
+    size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda"))
+    if local_size != max_size:
+        padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda")
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def mean(lst):
+    return sum(lst) / len(lst)
+
+
+def stop_gradient_by_name(name: str):
+    def apply_fn(module):
+        if hasattr(module, name):
+            getattr(module, name).requires_grad_(False)
+
+    return apply_fn

omnilmm/train/train_utils.py

@@ -0,0 +1,153 @@
+import os
+import gc
+import copy
+import time
+
+import torch
+import warnings
+import transformers
+
+import numpy as np
+
+from typing import Dict, Optional, Sequence
+from omnilmm import conversation as conversation_lib
+
+IGNORE_INDEX = -100
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+
+
+def _tokenize_fn(strings: Sequence[str],
+                 tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        ) for text in strings
+    ]
+    input_ids = labels = [
+        tokenized.input_ids[0] for tokenized in tokenized_list
+    ]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()
+        for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+
+
+def omni_preprocess(sources,
+                      tokenizer: transformers.PreTrainedTokenizer,
+                      generation=False):
+    system_content = 'You are an artificial intelligence assistant, which gives helpful, detailed, and polite answers to the human\'s questions.'
+    ignore_index = -100
+
+    response_template = '\n<|assistant|>\n'
+    instruction_template = '\n<|user|>\n'
+    response_token_ids = tokenizer.encode(
+        response_template, add_special_tokens=False)
+    instruction_token_ids = tokenizer.encode(
+        instruction_template, add_special_tokens=False)
+
+    batch_input_ids = []
+    batch_labels = []
+    for i in range(len(sources)):
+        new_source = []
+        prev_role = 'unexpect'
+        for conv_turn in sources[i]:
+            role = conv_turn['from'] if 'from' in conv_turn else conv_turn['role']
+            content = conv_turn['value'] if 'value' in conv_turn else conv_turn['content']
+
+            role = 'user' if role == 'human' else role
+            role = 'assistant' if role == 'gpt' else role
+
+            assert role in ['user', 'assistant']
+            assert role != prev_role, f'role={role}, prev_role={prev_role}'
+            prev_role = role
+
+            new_turn = {
+                'role': role,
+                'content': content
+            }
+            new_source.append(new_turn)
+        if new_source[0]['role'] != 'system':
+            new_source.insert(0, {'role': 'system', 'content': system_content})
+
+        # TODO: this automatically add '\n' to the end
+        res_text = tokenizer.apply_chat_template(
+            new_source, tokenize=False, add_generation_prompt=generation)
+        if not generation:
+            res_text = res_text.strip()
+
+        conversations_tokenized = _tokenize_fn([res_text], tokenizer)
+        res_input_ids = conversations_tokenized["input_ids"][0]
+
+        # since labels and input_ids are reference towards the same object
+        res_labels = copy.deepcopy(conversations_tokenized["labels"][0])
+
+        response_token_ids_idxs = []
+        human_token_ids_idxs = []
+
+        for assistant_idx in np.where(res_labels == response_token_ids[0])[0]:
+            # find the indexes of the start of a response.
+            if (response_token_ids == res_labels[assistant_idx: assistant_idx + len(
+                        response_token_ids)].tolist()
+                    ):
+                response_token_ids_idxs.append(
+                    assistant_idx + len(response_token_ids))
+
+        if len(response_token_ids_idxs) == 0:
+            warnings.warn(
+                f"Could not find response key `{response_template}` in the "
+                f'following instance: @===>{tokenizer.decode(res_input_ids)}<===@ '
+                f'Raw text is @===>{res_text}<===@'
+                f'Raw source is @===>{new_source}<===@'
+                f"This instance will be ignored in loss calculation. "
+                f"Note, if this happens often, consider increasing the `max_seq_length`."
+            )
+            res_labels[:] = ignore_index
+
+        human_token_ids = instruction_token_ids
+        for human_idx in np.where(res_labels == human_token_ids[0])[0]:
+            # find the indexes of the start of a human answer.
+            if human_token_ids == res_labels[human_idx: human_idx + len(human_token_ids)].tolist():
+                human_token_ids_idxs.append(human_idx)
+
+        if len(human_token_ids_idxs) == 0:
+            warnings.warn(
+                f"Could not find instruction key `{instruction_template}` in the "
+                f'following instance: @===>{tokenizer.decode(res_input_ids)}<===@ '
+                f'Raw text is @===>{res_text}<===@'
+                f'Raw source is @===>{new_source}<===@'
+                f"This instance will be ignored in loss calculation. "
+                f"Note, if this happens often, consider increasing the `max_seq_length`."
+            )
+            res_labels[:] = ignore_index
+
+        for idx, (start, end) in enumerate(zip(human_token_ids_idxs, response_token_ids_idxs)):
+            # Make pytorch loss function ignore all non response tokens
+            if idx != 0:
+                res_labels[start:end] = ignore_index
+            else:
+                res_labels[:end] = ignore_index
+
+        if len(response_token_ids_idxs) < len(human_token_ids_idxs):
+            res_labels[human_token_ids_idxs[-1]:] = ignore_index
+
+        batch_input_ids.append(res_input_ids)
+        batch_labels.append(res_labels)
+
+    return dict(input_ids=batch_input_ids, labels=batch_labels)
+
+

omnilmm/utils.py

@@ -0,0 +1,127 @@
+import datetime
+import logging
+import logging.handlers
+import os
+import sys
+
+import requests
+
+from omnilmm.constants import LOGDIR
+
+server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
+moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
+
+handler = None
+
+
+def build_logger(logger_name, logger_filename):
+    global handler
+
+    formatter = logging.Formatter(
+        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    # Set the format of root handlers
+    if not logging.getLogger().handlers:
+        logging.basicConfig(level=logging.INFO)
+    logging.getLogger().handlers[0].setFormatter(formatter)
+
+    # Redirect stdout and stderr to loggers
+    stdout_logger = logging.getLogger("stdout")
+    stdout_logger.setLevel(logging.INFO)
+    sl = StreamToLogger(stdout_logger, logging.INFO)
+    sys.stdout = sl
+
+    stderr_logger = logging.getLogger("stderr")
+    stderr_logger.setLevel(logging.ERROR)
+    sl = StreamToLogger(stderr_logger, logging.ERROR)
+    sys.stderr = sl
+
+    # Get logger
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+
+    # Add a file handler for all loggers
+    if handler is None:
+        os.makedirs(LOGDIR, exist_ok=True)
+        filename = os.path.join(LOGDIR, logger_filename)
+        handler = logging.handlers.TimedRotatingFileHandler(
+            filename, when='D', utc=True)
+        handler.setFormatter(formatter)
+
+        for name, item in logging.root.manager.loggerDict.items():
+            if isinstance(item, logging.Logger):
+                item.addHandler(handler)
+
+    return logger
+
+
+class StreamToLogger(object):
+    """
+    Fake file-like stream object that redirects writes to a logger instance.
+    """
+
+    def __init__(self, logger, log_level=logging.INFO):
+        self.terminal = sys.stdout
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ''
+
+    def __getattr__(self, attr):
+        return getattr(self.terminal, attr)
+
+    def write(self, buf):
+        temp_linebuf = self.linebuf + buf
+        self.linebuf = ''
+        for line in temp_linebuf.splitlines(True):
+            # From the io.TextIOWrapper docs:
+            #   On output, if newline is None, any '\n' characters written
+            #   are translated to the system default line separator.
+            # By default sys.stdout.write() expects '\n' newlines and then
+            # translates them so this is still cross platform.
+            if line[-1] == '\n':
+                self.logger.log(self.log_level, line.rstrip())
+            else:
+                self.linebuf += line
+
+    def flush(self):
+        if self.linebuf != '':
+            self.logger.log(self.log_level, self.linebuf.rstrip())
+        self.linebuf = ''
+
+
+def disable_torch_init():
+    """
+    Disable the redundant torch default initialization to accelerate model creation.
+    """
+    import torch
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+
+def violates_moderation(text):
+    """
+    Check whether the text violates OpenAI moderation API.
+    """
+    url = "https://api.openai.com/v1/moderations"
+    headers = {"Content-Type": "application/json",
+               "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"]}
+    text = text.replace("\n", "")
+    data = "{" + '"input": ' + f'"{text}"' + "}"
+    data = data.encode("utf-8")
+    try:
+        ret = requests.post(url, headers=headers, data=data, timeout=5)
+        flagged = ret.json()["results"][0]["flagged"]
+    except requests.exceptions.RequestException as e:
+        flagged = False
+    except KeyError as e:
+        flagged = False
+
+    return flagged
+
+
+def pretty_print_semaphore(semaphore):
+    if semaphore is None:
+        return "None"
+    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"

openai_api.py

@@ -0,0 +1,501 @@
+import gc
+import json
+import time
+import requests
+import base64
+import uvicorn
+import argparse
+
+import torch
+from transformers import AutoModelForCausalLM, LlamaTokenizer, PreTrainedModel, PreTrainedTokenizer, \
+TextIteratorStreamer, CodeGenTokenizerFast as Tokenizer
+
+from contextlib import asynccontextmanager
+from loguru import logger
+from typing import List, Literal, Union, Tuple, Optional
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+
+from PIL import Image
+from io import BytesIO
+
+import os
+import re
+from threading import Thread
+from moondream import Moondream, detect_device
+
+import omnichat
+
+# 请求
+class TextContent(BaseModel):
+    type: Literal["text"]
+    text: str
+class ImageUrl(BaseModel):
+    url: str
+class ImageUrlContent(BaseModel):
+    type: Literal["image_url"]
+    image_url: ImageUrl
+ContentItem = Union[TextContent, ImageUrlContent]
+class ChatMessageInput(BaseModel):
+    role: Literal["user", "assistant", "system"]
+    content: Union[str, List[ContentItem]]
+    name: Optional[str] = None
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessageInput]
+    temperature: Optional[float] = 0.8
+    top_p: Optional[float] = 0.8
+    max_tokens: Optional[int] = None
+    stream: Optional[bool] = False
+    # Additional parameters
+    repetition_penalty: Optional[float] = 1.0
+
+# 响应
+class ChatMessageResponse(BaseModel):
+    role: Literal["assistant"]
+    content: str = None
+    name: Optional[str] = None
+class ChatCompletionResponseChoice(BaseModel):
+    index: int
+    message: ChatMessageResponse
+class DeltaMessage(BaseModel):
+    role: Optional[Literal["user", "assistant", "system"]] = None
+    content: Optional[str] = None
+class ChatCompletionResponseStreamChoice(BaseModel):
+    index: int
+    delta: DeltaMessage
+class UsageInfo(BaseModel):
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    completion_tokens: Optional[int] = 0
+class ChatCompletionResponse(BaseModel):
+    model: str
+    object: Literal["chat.completion", "chat.completion.chunk"]
+    choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]]
+    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
+    usage: Optional[UsageInfo] = None
+
+# 图片输入处理
+def process_img(input_data):
+    if isinstance(input_data, str):
+        # URL
+        if input_data.startswith("http://") or input_data.startswith("https://"):
+            response = requests.get(input_data)
+            image_data = response.content
+            pil_image = Image.open(BytesIO(image_data)).convert('RGB')
+        # base64
+        elif input_data.startswith("data:image/"):
+            base64_data = input_data.split(",")[1]
+            image_data = base64.b64decode(base64_data)
+            pil_image = Image.open(BytesIO(image_data)).convert('RGB')
+        # img_path
+        else:
+            pil_image = Image.open(input_data)
+    # PIL
+    elif isinstance(input_data, Image.Image):
+        pil_image = input_data
+    else:
+        raise ValueError("data type error")
+
+    return pil_image
+
+# 历史消息处理
+def process_history_and_images(messages: List[ChatMessageInput]) -> Tuple[
+    Optional[str], Optional[List[Tuple[str, str]]], Optional[List[Image.Image]]]:
+    formatted_history = []
+    image_list = []
+    last_user_query = ''
+
+    for i, message in enumerate(messages):
+        role = message.role
+        content = message.content
+
+        if isinstance(content, list):  # text
+            text_content = ' '.join(item.text for item in content if isinstance(item, TextContent))
+        else:
+            text_content = content
+
+        if isinstance(content, list):  # image
+            for item in content:
+                if isinstance(item, ImageUrlContent):
+                    image_url = item.image_url.url
+                    image = process_img(image_url)
+                    image_list.append(image)
+
+        if role == 'user':
+            if i == len(messages) - 1:  # last message
+                last_user_query = text_content
+            else:
+                formatted_history.append((text_content, ''))
+        elif role == 'assistant':
+            if formatted_history:
+                if formatted_history[-1][1] != '':
+                    assert False, f"the last query is answered. answer again. {formatted_history[-1][0]}, {formatted_history[-1][1]}, {text_content}"
+                formatted_history[-1] = (formatted_history[-1][0], text_content)
+            else:
+                assert False, f"assistant reply before user"
+        else:
+            assert False, f"unrecognized role: {role}"
+
+    return last_user_query, formatted_history, image_list
+
+@torch.inference_mode()
+# Moondrean推理
+def generate_stream_moondream(params: dict):
+    global model, tokenizer
+
+    # 输入处理
+    def chat_history_to_prompt(history):
+        prompt = ""
+        for i, (old_query, response) in enumerate(history):
+            prompt += f"Question: {old_query}\n\nAnswer: {response}\n\n"
+        return prompt
+
+    messages = params["messages"]
+    prompt, formatted_history, image_list = process_history_and_images(messages)
+    history = chat_history_to_prompt(formatted_history)
+    # 只处理最后一张图
+    img = image_list[-1]
+
+    # 构建输入
+    '''
+    answer_question(
+            self,
+            image_embeds,
+            question,
+            tokenizer,
+            chat_history="",
+            result_queue=None,
+            **kwargs,
+        )
+    '''
+    image_embeds = model.encode_image(img)
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    gen_kwargs = {
+            "image_embeds": image_embeds,
+            "question": prompt,
+            "tokenizer": tokenizer,
+            "chat_history": history,
+            "result_queue": None,
+            "streamer": streamer,
+        }
+
+    thread = Thread(
+        target=model.answer_question,
+        kwargs=gen_kwargs,
+    )
+
+    input_echo_len = 0
+    total_len = 0
+    # 启动推理
+    thread.start()
+    buffer = ""
+    for new_text in streamer:
+        clean_text = re.sub("<$|END$", "", new_text)
+        buffer += clean_text
+        yield {
+            "text": buffer.strip("<END"),
+            "usage": {
+                "prompt_tokens": input_echo_len,
+                "completion_tokens": total_len - input_echo_len,
+                "total_tokens": total_len,
+                },
+            }
+    generated_ret ={
+        "text": buffer.strip("<END"),
+        "usage": {
+            "prompt_tokens": input_echo_len,
+            "completion_tokens": total_len - input_echo_len,
+            "total_tokens": total_len,
+            },
+        }
+    yield generated_ret
+
+# Moondrean单次响应
+def generate_moondream(params: dict):
+    for response in generate_stream_moondream(params):
+        pass
+    return response
+
+
+@torch.inference_mode()
+# CogVLM推理
+def generate_stream_cogvlm(model: PreTrainedModel, tokenizer: PreTrainedTokenizer, params: dict):
+    """
+    Generates a stream of responses using the CogVLM model in inference mode.
+    It's optimized to handle continuous input-output interactions with the model in a streaming manner.
+    """
+    messages = params["messages"]
+    temperature = float(params.get("temperature", 1.0))
+    repetition_penalty = float(params.get("repetition_penalty", 1.0))
+    top_p = float(params.get("top_p", 1.0))
+    max_new_tokens = int(params.get("max_tokens", 256))
+    query, history, image_list = process_history_and_images(messages)
+
+    logger.debug(f"==== request ====\n{query}")
+
+    #  only can slove the latest picture
+    input_by_model = model.build_conversation_input_ids(tokenizer, query=query, history=history,
+                                                        images=[image_list[-1]])
+    inputs = {
+        'input_ids': input_by_model['input_ids'].unsqueeze(0).to(DEVICE),
+        'token_type_ids': input_by_model['token_type_ids'].unsqueeze(0).to(DEVICE),
+        'attention_mask': input_by_model['attention_mask'].unsqueeze(0).to(DEVICE),
+        'images': [[input_by_model['images'][0].to(DEVICE).to(torch_type)]],
+    }
+    if 'cross_images' in input_by_model and input_by_model['cross_images']:
+        inputs['cross_images'] = [[input_by_model['cross_images'][0].to(DEVICE).to(torch_type)]]
+
+    input_echo_len = len(inputs["input_ids"][0])
+    streamer = TextIteratorStreamer(tokenizer=tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs = {
+        "repetition_penalty": repetition_penalty,
+        "max_new_tokens": max_new_tokens,
+        "do_sample": False,
+        "top_p": top_p,
+        'streamer': streamer,
+    }
+    if temperature > 1e-5:
+        gen_kwargs["temperature"] = temperature
+
+    total_len = 0
+    generated_text = ""
+    with torch.no_grad():
+        model.generate(**inputs, **gen_kwargs)
+        for next_text in streamer:
+            generated_text += next_text
+            yield {
+                "text": generated_text,
+                "usage": {
+                    "prompt_tokens": input_echo_len,
+                    "completion_tokens": total_len - input_echo_len,
+                    "total_tokens": total_len,
+                },
+            }
+    ret = {
+        "text": generated_text,
+        "usage": {
+            "prompt_tokens": input_echo_len,
+            "completion_tokens": total_len - input_echo_len,
+            "total_tokens": total_len,
+        },
+    }
+    yield ret
+
+# CogVLM单次响应
+def generate_cogvlm(model: PreTrainedModel, tokenizer: PreTrainedTokenizer, params: dict):
+
+    for response in generate_stream_cogvlm(model, tokenizer, params):
+        pass
+    return response
+
+def generate_minicpm(model, params):
+    messages = params["messages"]
+    query, history, image_list = process_history_and_images(messages)
+    msgs = history
+    msgs.append({'role': 'user', 'content': query})
+    image = image_list[-1]
+    # image is a PIL image
+    buffer = BytesIO()
+    image.save(buffer, format="JPEG")  # You can adjust the format as needed
+    buffer.seek(0)
+    image_base64 = base64.b64encode(buffer.read())
+    image_base64_str = image_base64.decode("utf-8")
+    input = {'image': image_base64_str, 'question': json.dumps(msgs)}
+    generation = model.chat(input)
+    response = {"text": generation, "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}}
+    print(response)
+    return response
+
+# 流式响应
+async def predict(model_id: str, params: dict):
+    return "no stream"
+
+torch.set_grad_enabled(False)
+# 生命周期管理器，结束清显存
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    yield
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+app = FastAPI(lifespan=lifespan)
+# 允许跨域
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 对话路由
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def create_chat_completion(request: ChatCompletionRequest):
+    global model, tokenizer
+
+    # 检查请求
+    if len(request.messages) < 1 or request.messages[-1].role == "assistant":
+        raise HTTPException(status_code=400, detail="Invalid request")
+
+    gen_params = dict(
+        messages=request.messages,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        max_tokens=request.max_tokens or 1024,
+        echo=False,
+        stream=request.stream,
+    )
+
+    # 流式响应
+    if request.stream:
+        generate = predict(request.model, gen_params)
+        return
+
+    # 单次响应
+    if STATE_MOD == "cog":
+        response = generate_cogvlm(model, tokenizer, gen_params)
+    elif STATE_MOD == "moon":
+        response = generate_moondream(gen_params)
+    elif STATE_MOD == "mini":
+        response = generate_minicpm(model, gen_params)
+    usage = UsageInfo()
+    message = ChatMessageResponse(
+        role="assistant",
+        content=response["text"],
+    )
+    logger.debug(f"==== message ====\n{message}")
+    choice_data = ChatCompletionResponseChoice(
+        index=0,
+        message=message,
+    )
+    task_usage = UsageInfo.model_validate(response["usage"])
+    for usage_key, usage_value in task_usage.model_dump().items():
+        setattr(usage, usage_key, getattr(usage, usage_key) + usage_value)
+    return ChatCompletionResponse(model=request.model, choices=[choice_data], object="chat.completion", usage=usage)
+
+# 模型切换路由配置
+STATE_MOD = "moon"
+MODEL_PATH = ""
+
+# 模型加载
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+def load_mod(model_input, mod_type):
+    global model, tokenizer, language_processor_version
+    if mod_type == "cog":
+        tokenizer_path = os.environ.get("TOKENIZER_PATH", 'lmsys/vicuna-7b-v1.5')
+        tokenizer = LlamaTokenizer.from_pretrained(
+            tokenizer_path,
+            trust_remote_code=True,
+            signal_type=language_processor_version
+        )
+        if 'cuda' in DEVICE:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_input,
+                trust_remote_code=True,
+                load_in_4bit=True,
+                torch_dtype=torch_type,
+                low_cpu_mem_usage=True
+            ).eval()
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_input,
+                trust_remote_code=True
+            ).float().to(DEVICE).eval()
+    elif mod_type == "moon":
+        device, dtype = detect_device()
+        model = Moondream.from_pretrained(model_input).to(device=device, dtype=dtype).eval()
+        tokenizer = Tokenizer.from_pretrained(model_input)
+    elif mod_type == "mini":
+        model, tokenizer = omnichat.OmniLMMChat(model_input), None
+
+@app.post("/v1/Cog-vqa")
+async def switch_vqa():
+    global model, STATE_MOD, mod_vqa, language_processor_version
+    STATE_MOD = "cog"
+    del model
+    model = None
+    language_processor_version = "chat_old"
+    load_mod(mod_vqa, STATE_MOD)
+
+@app.post("/v1/Cog-chat")
+async def switch_chat():
+    global model, STATE_MOD, mod_chat, language_processor_version
+    STATE_MOD = "cog"
+    del model
+    model = None
+    language_processor_version = "chat"
+    load_mod(mod_chat, STATE_MOD)
+
+@app.post("/v1/moondream")
+async def switch_moon():
+    global model, STATE_MOD, mod_moon
+    STATE_MOD = "moon"
+    del model
+    model = None
+    load_mod(mod_moon, STATE_MOD)
+
+@app.post("/v1/MiniCPM")
+async def switch_mini():
+    global model, STATE_MOD, mod_mini
+    STATE_MOD = "mini"
+    del model
+    model = None
+    load_mod(mod_mini, STATE_MOD)
+
+# 关闭
+@app.post("/v1/close")
+async def close():
+    global model
+    del model
+    model = None
+
+gc.collect()
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--mod", type=str, default="moondrean")
+args = parser.parse_args()
+mod = args.mod
+
+mod_vqa = './models/cogagent-vqa-hf'
+mod_chat = './models/cogagent-chat-hf'
+mod_moon = './models/moondream'
+mod_mini = './models/MiniCPM-Llama3-V-2_5'
+
+'''
+mod_list = [
+    "moondrean",
+    "Cog-vqa",
+    "Cog-chat"
+    "MiniCPM"
+    ]
+'''
+
+if mod == "Cog-vqa":
+    STATE_MOD = "cog"
+    MODEL_PATH = mod_vqa
+    language_processor_version = "chat_old"
+elif mod == "Cog-chat":
+    STATE_MOD = "cog"
+    MODEL_PATH = mod_chat
+    language_processor_version = "chat"
+elif mod == "moondream":
+    STATE_MOD = "moon"
+    MODEL_PATH = mod_moon
+elif mod == "MiniCPM":
+    STATE_MOD = "mini"
+    MODEL_PATH = mod_mini
+
+if __name__ == "__main__":
+    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
+        torch_type = torch.bfloat16
+    else:
+        torch_type = torch.float16
+
+    print("========Use torch type as:{} with device:{}========\n\n".format(torch_type, DEVICE))
+
+    load_mod(MODEL_PATH, STATE_MOD)
+
+    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)

start_linux_mac.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+export HF_HOME="huggingface"
+
+python ./install_script/check_open.py
+
+python gpt-caption.py --share "$@"
+
+read -p "Press any key to continue . . . "

start_windows.bat

@@ -0,0 +1,9 @@
+@echo off
+set HF_HOME=huggingface
+
+call myenv\Scripts\activate
+python ./install_script/check_open.py
+
+python gpt-caption.py %*
+
+pause

utils/merge_model.py

@@ -0,0 +1,42 @@
+# -*- encoding: utf-8 -*-
+import os, sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch
+import argparse
+from models.cogvlm_model import FineTuneTestCogVLMModel
+from sat.training.model_io import save_checkpoint
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--version", type=str, default="base", help='version to interact with')
+    parser.add_argument("--from_pretrained", type=str, default="checkpoints/merged_lora", help='pretrained ckpt')
+    parser.add_argument("--fp16", action="store_true")
+    parser.add_argument("--bf16", action="store_true")
+    args = parser.parse_args()
+    rank = int(os.environ.get('RANK', 0))
+    world_size = int(os.environ.get('WORLD_SIZE', 1))
+    parser = FineTuneTestCogVLMModel.add_model_specific_args(parser)
+    args = parser.parse_args()
+
+    # load model
+    model, model_args = FineTuneTestCogVLMModel.from_pretrained(
+        args.from_pretrained,
+        args=argparse.Namespace(
+        deepspeed=None,
+        local_rank=rank,
+        rank=rank,
+        world_size=world_size,
+        model_parallel_size=world_size,
+        mode='inference',
+        skip_init=True,
+        use_gpu_initialization=True if torch.cuda.is_available() else False,
+        device='cuda',
+        **vars(args)
+    ), url='local', overwrite_args={'model_parallel_size': 1})
+    model = model.eval()
+    model_args.save = './checkpoints/merged_model_{}'.format(model_args.eva_args["image_size"][0])
+    save_checkpoint(1, model, None, None, model_args)
+
+if __name__ == "__main__":
+    main()

utils/split_dataset.py

@@ -0,0 +1,35 @@
+import os
+import shutil
+
+def find_all_files(path, suffix=".jpg"):
+    target_files = []
+    for cur_dir, _, files in os.walk(path, followlinks=True):
+        for f in files:
+            if f.endswith(suffix):
+                target_files.append(os.path.join(cur_dir, f))
+    print(f'find {len(target_files)} files...')
+    return target_files
+
+all_files = find_all_files('archive')
+os.makedirs("archive_split", exist_ok=True)
+os.makedirs("archive_split/train", exist_ok=True)
+os.makedirs("archive_split/valid", exist_ok=True)
+os.makedirs("archive_split/test", exist_ok=True)
+
+import random
+random.seed(2023)
+random.shuffle(all_files)
+train = all_files[:8000]
+valid = all_files[8000:8000+500]
+test = all_files[8000+500:8000+500+1500]
+
+print("building train")
+for file in train:
+    shutil.move(file, os.path.join("archive_split/train", file.split("/")[-1]))
+print("building valid")
+for file in valid:
+    shutil.move(file, os.path.join("archive_split/valid", file.split("/")[-1]))
+print("building test")
+for file in test:
+    shutil.move(file, os.path.join("archive_split/test", file.split("/")[-1]))
+print("done")

utils/utils/__init__.py

@@ -0,0 +1,5 @@
+from .chat import chat
+from .language import llama2_tokenizer, llama2_text_processor, llama2_text_processor_inference
+from .vision import get_image_processor
+from .grounding_parser import parse_response
+from .dataset import ItemDataset

utils/utils/chat.py

@@ -0,0 +1,149 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   chat.py
+@Time    :   2023/05/08 19:10:08
+@Author  :   Ming Ding 
+@Contact :   dm18@mails.tsinghua.edu.cn
+'''
+
+from typing import Optional, Tuple, Union, List, Callable, Dict, Any
+import requests
+from PIL import Image
+from io import BytesIO
+
+import torch
+from sat.generation.autoregressive_sampling import filling_sequence, stream_filling_sequence, get_masks_and_position_ids_default
+from sat.generation.sampling_strategies import BaseStrategy, BeamSearchStrategy
+from sat.mpu import get_model_parallel_rank
+
+def process_image(image_path, img_processor, cross_img_processor, image):
+    if image is None:
+        if image_path.startswith("http"):
+            response = requests.get(image_path, timeout=10)
+            image = Image.open(BytesIO(response.content))
+        else:
+            image = Image.open(image_path)
+
+    if image is not None and isinstance(image, Image.Image):
+        pil_img = image.convert('RGB')
+        img_dict = img_processor(pil_img)
+        cross_img_dict = cross_img_processor(pil_img) if cross_img_processor is not None else {}
+        ret = (img_dict, pil_img, cross_img_dict)
+    else:
+        ret = image
+    return ret
+
+def chat(image_path, model, text_processor, img_processor,
+        query: str, history: List[Tuple[str, str]] = None, cross_img_processor=None, image: Image = None,
+        max_length: int = 4096, top_p=0.95, top_k=5, temperature=0.95, repetition_penalty=1.0,
+        invalid_slices=[], no_prompt=False, args=None
+        ):
+    if image is None:
+        assert image_path is not None
+    if not history:
+        history = []
+
+    if no_prompt:
+        query = ''
+    prompt = text_processor.history_to_prompt(query, history)
+
+    (torch_image, pil_img, cross_image) = process_image(image_path, img_processor, cross_img_processor, image)
+
+    if torch_image is not None:
+        for k in torch_image:
+            if type(torch_image[k]) is torch.Tensor and torch_image[k].dtype is not torch.int and torch_image[k].dtype is not torch.long:
+                torch_image[k] = torch_image[k].to(torch.bfloat16 if args.bf16 else torch.float16)
+            if type(torch_image[k]) is torch.Tensor:
+                torch_image[k] = torch_image[k].to(next(model.parameters()).device)
+                
+    if cross_image is not None:
+        for k in cross_image:
+            if type(cross_image[k]) is torch.Tensor and cross_image[k].dtype is not torch.int and cross_image[k].dtype is not torch.long:
+                cross_image[k] = cross_image[k].to(torch.bfloat16 if args.bf16 else torch.float16)
+            if type(cross_image[k]) is torch.Tensor:
+                cross_image[k] = cross_image[k].to(next(model.parameters()).device)
+
+    inputs_dic = text_processor(prompt)
+    for k in inputs_dic:
+        if type(inputs_dic[k]) is torch.Tensor and inputs_dic[k].dtype is not torch.int and inputs_dic[k].dtype is not torch.long:
+            inputs_dic[k] = inputs_dic[k].to(torch.bfloat16 if args.bf16 else torch.float16)
+        if type(inputs_dic[k]) is torch.Tensor:
+            inputs_dic[k] = inputs_dic[k].to(next(model.parameters()).device)
+    input_ids = inputs_dic['input_ids'].to(model.parameters().__next__().device)[0]
+    
+    if max_length-len(input_ids) <= 1:
+        response = "The prompt exceeds the context length limit, please try again."
+        return response, history, (torch_image, pil_img)
+    
+    seq = torch.cat(
+        [input_ids, torch.tensor([-1]*(max_length-len(input_ids)), device=input_ids.device)], dim=0
+    )
+    strategy = BaseStrategy(temperature=temperature, top_p=top_p, top_k=top_k, end_tokens=[text_processor.tokenizer.eos_token_id],
+                            invalid_slices=invalid_slices, repetition_penalty=repetition_penalty)
+    # use beam search to get a better result
+    # strategy = BeamSearchStrategy(temperature=temperature, top_p=top_p, top_k=top_k, end_tokens=[text_processor.tokenizer.eos_token_id],
+    #                               num_beams=5, consider_end=True, repetition_penalty=repetition_penalty)
+    get_func = text_processor.get_func(input_ids, **inputs_dic) if hasattr(text_processor, 'get_func') else get_masks_and_position_ids_default
+
+    img_inputs = {'vision_'+k: v for k, v in torch_image.items()}
+    if cross_image is not None:
+        img_inputs = {**img_inputs, **{'cross_'+k:v for k,v in cross_image.items()}}
+    inputs_dic.pop('input_ids')
+    inputs = {**img_inputs, **inputs_dic}
+
+    if args.stream_chat:
+        filling_stream = stream_filling_sequence(
+            model, seq,
+            batch_size=1,
+            get_masks_and_position_ids=get_func,
+            strategy=strategy,
+            **inputs
+        )
+        if get_model_parallel_rank() == 0:
+            if 'chinese' in args and not args.chinese:
+                print("Model: ", end='')
+            else:
+                print("模型��", end='')
+        offset = len(text_processor.tokenizer.decode(input_ids))
+        for tokens, mems in filling_stream:
+            torch.cuda.empty_cache()
+            tmp_response = text_processor.tokenizer.decode(tokens[0])
+            if tmp_response[-1] != "�":
+                if get_model_parallel_rank() == 0:
+                    tmp_response_offseted = tmp_response[offset:]
+                    if hasattr(text_processor, 'process_response'):
+                        tmp_response_offseted = text_processor.process_response(tmp_response_offseted)
+                    print(tmp_response_offseted, end='', flush=True)
+                offset = len(tmp_response)
+        if get_model_parallel_rank() == 0:
+            print()
+        output = strategy.finalize(tokens, mems)[0]
+
+        response = text_processor.tokenizer.decode(output[0])
+    else:
+        output = filling_sequence(
+            model, seq,
+            batch_size=1,
+            get_masks_and_position_ids=get_func,
+            strategy=strategy,
+            **inputs
+        )[0] # drop memory
+        
+        # ---------------
+        # port from inference_glm.py, more general than chat mode
+        # clip -1s and fill back generated things into seq
+        if type(output) is not list:
+            output_list = output.tolist()
+        else:
+            output_list = output
+
+        response = text_processor.tokenizer.decode(output_list[0])
+    # print('original:', response)
+    if hasattr(text_processor, 'process_response'):
+        response = text_processor.process_response(response)
+    response = response.split(text_processor.sep)[-1].strip()
+    if get_model_parallel_rank() == 0:
+        from utils.utils.grounding_parser import parse_response
+        parse_response(pil_img, response)
+    history = history + [(query, response)]
+    return response, history, (torch_image, pil_img, cross_image)

utils/utils/dataset.py

@@ -0,0 +1,61 @@
+import os
+import logging
+import random
+import logging
+import jsonlines
+from io import BytesIO
+from PIL import Image
+from torch.utils.data import Dataset
+from sat.helpers import print_rank0
+
+def find_all_files(path, suffix=".jpg"):
+    target_files = []
+    for cur_dir, _, files in os.walk(path, followlinks=True):
+        for f in files:
+            if f.endswith(suffix):
+                target_files.append(os.path.join(cur_dir, f))
+    print_rank0(f'find {len(target_files)} files...')
+    return target_files
+
+class ItemDataset(Dataset):
+    def __init__(self, image_processor, text_processor, args, data_dirs, cross_image_processor=None, **kwargs):
+        super().__init__()
+        self.data = self.load_data(data_dirs)
+        self.image_processor, self.text_processor, self.cross_image_processor = image_processor, text_processor, cross_image_processor
+    
+    def process_img(self, img):
+        img_dict = {'vision': self.image_processor(img)}
+        if self.cross_image_processor:
+            img_dict.update({'cross': self.cross_image_processor(img)})
+        return img_dict
+    
+    def process_text(self, answer, prompt):
+        return self.text_processor(answer, prompt)
+    
+    def load_data(self, data_dir):
+        all_files = find_all_files(data_dir, suffix=".jpg")
+        print_rank0(f"find {len(all_files)} samples in all...")
+        return all_files
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        data = self.data[index]
+        # img
+        try:
+            img = Image.open(data).convert('RGB')
+        except Exception as e:
+            print_rank0(e, level=logging.WARNING)
+            return {}
+        img_dict = self.process_img(img)
+        # text
+        label = data.split('/')[-1].split('.')[0]
+        uni_key = label
+        text_dict = self.process_text(label, "CAPTCHA:")
+        if text_dict is None:
+            print_rank0(f"Process text failed. Please check the max_target_length & max_source_length.\n The data is {data}", level=logging.WARNING)
+            return {}
+        # other attr
+        ret = {**img_dict, **text_dict, "question_id": uni_key}
+        return ret

utils/utils/grounding_parser.py

@@ -0,0 +1,86 @@
+import seaborn as sns
+from PIL import Image, ImageDraw, ImageFont
+import matplotlib.font_manager
+import spacy
+import re
+
+nlp = spacy.load("en_core_web_sm")
+
+def draw_boxes(image, boxes, texts, output_fn='output.png'):
+    box_width = 5
+    color_palette = sns.color_palette("husl", len(boxes))
+    colors = [(int(r*255), int(g*255), int(b*255)) for r, g, b in color_palette]
+
+    width, height = image.size
+    absolute_boxes = [[(int(box[0] * width), int(box[1] * height), int(box[2] * width), int(box[3] * height)) for box in b] for b in boxes]
+    
+    overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
+    draw = ImageDraw.Draw(overlay)
+    font_path = sorted(matplotlib.font_manager.findSystemFonts(fontpaths=None, fontext='ttf'))[0]
+    font = ImageFont.truetype(font_path, size=26)
+
+    for box, text, color in zip(absolute_boxes, texts, colors):
+        for b in box:
+            draw.rectangle(b, outline=color, width=box_width)
+            if not text:
+                continue
+            splited_text = text.split('\n')
+            num_lines = len(splited_text)
+            text_width, text_height = font.getbbox(splited_text[0])[-2:]
+            y_start = b[3] - text_height * num_lines - box_width
+            if b[2] - b[0] < 100 or b[3] - b[1] < 100:
+                y_start = b[3]
+            for i, line in enumerate(splited_text):
+                text_width, text_height = font.getbbox(line)[-2:]
+                x = b[0] + box_width
+                y = y_start + text_height * i
+                draw.rectangle([x, y, x+text_width, y+text_height], fill=(128, 128, 128, 160))
+                draw.text((x, y), line, font=font, fill=(255, 255, 255))
+    img_with_overlay = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
+    img_with_overlay.save(output_fn)
+
+def boxstr_to_boxes(box_str):
+    boxes = [[int(y)/1000 for y in x.split(',')] for x in box_str.split(';') if x.replace(',', '').isdigit()]
+    return boxes
+
+def text_to_dict(text):
+    doc = nlp(text)
+
+    box_matches = list(re.finditer(r'\[\[([^\]]+)\]\]', text))
+    box_positions = [match.start() for match in box_matches]
+
+    noun_phrases = []
+    boxes = []
+
+    for match, box_position in zip(box_matches, box_positions):
+        nearest_np_start = max([0] + [chunk.start_char for chunk in doc.noun_chunks if chunk.end_char <= box_position])
+        noun_phrase = text[nearest_np_start:box_position].strip()
+        if noun_phrase and noun_phrase[-1] == '?':
+            noun_phrase = text[:box_position].strip()
+        box_string = match.group(1)
+        
+        noun_phrases.append(noun_phrase)
+        boxes.append(boxstr_to_boxes(box_string))
+
+    pairs = []
+    for noun_phrase, box_string in zip(noun_phrases, boxes):
+        pairs.append((noun_phrase.lower(), box_string))
+    return dict(pairs)
+
+def parse_response(img, response, output_fn='output.png'):
+    img = img.convert('RGB')
+    width, height = img.size
+    ratio = min(1920 / width, 1080 / height)
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+    new_img = img.resize((new_width, new_height), Image.LANCZOS)
+    pattern = r"\[\[(.*?)\]\]"
+    positions = re.findall(pattern, response)
+    boxes = [[[int(y) for y in x.split(',')] for x in pos.split(';') if x.replace(',', '').isdigit()] for pos in positions]
+    dic = text_to_dict(response)
+    if not dic:
+        texts = []
+        boxes = []
+    else:
+        texts, boxes = zip(*dic.items())
+    draw_boxes(new_img, boxes, texts, output_fn=output_fn)