joycaption.py

import spaces
import gradio as gr
from huggingface_hub import InferenceClient
from torch import nn
from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
from pathlib import Path
import torch
import torch.amp.autocast_mode
from PIL import Image
import os
import torchvision.transforms.functional as TVF
import gc
from peft import PeftConfig

device = "cuda" if torch.cuda.is_available() else "cpu"
HF_TOKEN = os.environ.get("HF_TOKEN", None)
use_inference_client = False

llm_models = {
    "bunnycore/LLama-3.1-8B-Matrix": None,
    "Sao10K/Llama-3.1-8B-Stheno-v3.4": None,
    "unsloth/Meta-Llama-3.1-8B-bnb-4bit": None,
    "DevQuasar/HermesNova-Llama-3.1-8B": None,
    "mergekit-community/L3.1-Boshima-b-FIX": None,
    "meta-llama/Meta-Llama-3.1-8B": None, # gated
}

CLIP_PATH = "google/siglip-so400m-patch14-384"
MODEL_PATH = list(llm_models.keys())[0]
CHECKPOINT_PATH = Path("9em124t2-499968")
LORA_PATH = CHECKPOINT_PATH / "text_model"
TITLE = "<h1><center>JoyCaption Alpha One (2024-09-20a)</center></h1>"
CAPTION_TYPE_MAP = {
    ("descriptive", "formal", False, False): ["Write a descriptive caption for this image in a formal tone."],
    ("descriptive", "formal", False, True): ["Write a descriptive caption for this image in a formal tone within {word_count} words."],
    ("descriptive", "formal", True, False): ["Write a {length} descriptive caption for this image in a formal tone."],
    ("descriptive", "informal", False, False): ["Write a descriptive caption for this image in a casual tone."],
    ("descriptive", "informal", False, True): ["Write a descriptive caption for this image in a casual tone within {word_count} words."],
    ("descriptive", "informal", True, False): ["Write a {length} descriptive caption for this image in a casual tone."],

    ("training_prompt", "formal", False, False): ["Write a stable diffusion prompt for this image."],
    ("training_prompt", "formal", False, True): ["Write a stable diffusion prompt for this image within {word_count} words."],
    ("training_prompt", "formal", True, False): ["Write a {length} stable diffusion prompt for this image."],

    ("rng-tags", "formal", False, False): ["Write a list of Booru tags for this image."],
    ("rng-tags", "formal", False, True): ["Write a list of Booru tags for this image within {word_count} words."],
    ("rng-tags", "formal", True, False): ["Write a {length} list of Booru tags for this image."],
}

class ImageAdapter(nn.Module):
    def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
        super().__init__()
        self.deep_extract = deep_extract

        if self.deep_extract:
            input_features = input_features * 5

        self.linear1 = nn.Linear(input_features, output_features)
        self.activation = nn.GELU()
        self.linear2 = nn.Linear(output_features, output_features)
        self.ln1 = nn.Identity() if not ln1 else nn.LayerNorm(input_features)
        self.pos_emb = None if not pos_emb else nn.Parameter(torch.zeros(num_image_tokens, input_features))

        # Mode token
        #self.mode_token = nn.Embedding(n_modes, output_features)
        #self.mode_token.weight.data.normal_(mean=0.0, std=0.02)   # Matches HF's implementation of llama3

        # Other tokens (<|image_start|>, <|image_end|>, <|eot_id|>)
        self.other_tokens = nn.Embedding(3, output_features)
        self.other_tokens.weight.data.normal_(mean=0.0, std=0.02)   # Matches HF's implementation of llama3

    def forward(self, vision_outputs: torch.Tensor):
        if self.deep_extract:
            x = torch.concat((
                vision_outputs[-2],
                vision_outputs[3],
                vision_outputs[7],
                vision_outputs[13],
                vision_outputs[20],
            ), dim=-1)
            assert len(x.shape) == 3, f"Expected 3, got {len(x.shape)}"  # batch, tokens, features
            assert x.shape[-1] == vision_outputs[-2].shape[-1] * 5, f"Expected {vision_outputs[-2].shape[-1] * 5}, got {x.shape[-1]}"
        else:
            x = vision_outputs[-2]

        x = self.ln1(x)

        if self.pos_emb is not None:
            assert x.shape[-2:] == self.pos_emb.shape, f"Expected {self.pos_emb.shape}, got {x.shape[-2:]}"
            x = x + self.pos_emb

        x = self.linear1(x)
        x = self.activation(x)
        x = self.linear2(x)

        # Mode token
        #mode_token = self.mode_token(mode)
        #assert mode_token.shape == (x.shape[0], mode_token.shape[1], x.shape[2]), f"Expected {(x.shape[0], 1, x.shape[2])}, got {mode_token.shape}"
        #x = torch.cat((x, mode_token), dim=1)

        # <|image_start|>, IMAGE, <|image_end|>
        other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1))
        assert other_tokens.shape == (x.shape[0], 2, x.shape[2]), f"Expected {(x.shape[0], 2, x.shape[2])}, got {other_tokens.shape}"
        x = torch.cat((other_tokens[:, 0:1], x, other_tokens[:, 1:2]), dim=1)

        return x

    def get_eot_embedding(self):
        return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0)


# https://huggingface.co/docs/transformers/v4.44.2/gguf
# https://github.com/city96/ComfyUI-GGUF/issues/7
# https://github.com/THUDM/ChatGLM-6B/issues/18
# https://github.com/meta-llama/llama/issues/394
# https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct/discussions/109
# https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu
# https://huggingface.co/google/flan-ul2/discussions/8
# https://huggingface.co/blog/4bit-transformers-bitsandbytes
# https://huggingface.co/docs/transformers/main/en/peft
# https://huggingface.co/docs/transformers/main/en/peft#enable-and-disable-adapters
tokenizer = None
text_model_client = None
text_model = None
image_adapter = None
peft_config = None
def load_text_model(model_name: str=MODEL_PATH, gguf_file: str | None=None, is_nf4: bool=True):
    global tokenizer
    global text_model
    global image_adapter
    global peft_config
    global text_model_client #
    global use_inference_client #
    try:
        from transformers import BitsAndBytesConfig
        nf4_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                        bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
        print("Loading tokenizer")
        if gguf_file: tokenizer = AutoTokenizer.from_pretrained(model_name, gguf_file=gguf_file, use_fast=True, legacy=False)
        else: tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, legacy=False)
        assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
        print(f"Loading LLM: {model_name}")
        if gguf_file:
            if device == "cpu": text_model = AutoModelForCausalLM.from_pretrained(model_name, gguf_file=gguf_file, device_map=device, torch_dtype=torch.bfloat16).eval()
            elif is_nf4: text_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval()
            else: text_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval()
        else:
            if device == "cpu": text_model = AutoModelForCausalLM.from_pretrained(model_name, gguf_file=gguf_file, device_map=device, torch_dtype=torch.bfloat16).eval()
            elif is_nf4: text_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval()
            else: text_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval()
        if LORA_PATH.exists():
            print("Loading VLM's custom text model")
            if is_nf4: peft_config = PeftConfig.from_pretrained(LORA_PATH, device_map=device, quantization_config=nf4_config)
            else: peft_config = PeftConfig.from_pretrained(LORA_PATH, device_map=device)
            text_model.add_adapter(peft_config)
            text_model.enable_adapters()
        print("Loading image adapter")
        image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False).eval().to("cpu")
        image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu", weights_only=True))
        image_adapter.eval().to(device)
    except Exception as e:
        print(f"LLM load error: {e}")
        raise Exception(f"LLM load error: {e}") from e
    finally:
        torch.cuda.empty_cache()
        gc.collect()

load_text_model.zerogpu = True

# Load CLIP
print("Loading CLIP")
clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
clip_model = AutoModel.from_pretrained(CLIP_PATH).vision_model

if (CHECKPOINT_PATH / "clip_model.pt").exists():
    print("Loading VLM's custom vision model")
    checkpoint = torch.load(CHECKPOINT_PATH / "clip_model.pt", map_location='cpu')
    checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()}
    clip_model.load_state_dict(checkpoint)
    del checkpoint

clip_model.eval().requires_grad_(False).to(device)

# Tokenizer
# LLM
# Image Adapter
load_text_model()


@spaces.GPU()
@torch.no_grad()
def stream_chat(input_image: Image.Image, caption_type: str, caption_tone: str, caption_length: str | int) -> str:
    torch.cuda.empty_cache()

    # 'any' means no length specified
    length = None if caption_length == "any" else caption_length

    if isinstance(length, str):
        try:
            length = int(length)
        except ValueError:
            pass

    # 'rng-tags' and 'training_prompt' don't have formal/informal tones
    if caption_type == "rng-tags" or caption_type == "training_prompt":
        caption_tone = "formal"

    # Build prompt
    prompt_key = (caption_type, caption_tone, isinstance(length, str), isinstance(length, int))
    if prompt_key not in CAPTION_TYPE_MAP:
        raise ValueError(f"Invalid caption type: {prompt_key}")

    prompt_str = CAPTION_TYPE_MAP[prompt_key][0].format(length=length, word_count=length)
    print(f"Prompt: {prompt_str}")

    # Preprocess image
    #image = clip_processor(images=input_image, return_tensors='pt').pixel_values
    image = input_image.resize((384, 384), Image.LANCZOS)
    pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
    pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
    pixel_values = pixel_values.to('cuda')

    # Tokenize the prompt
    prompt = tokenizer.encode(prompt_str, return_tensors='pt', padding=False, truncation=False, add_special_tokens=False)

    # Embed image
    with torch.amp.autocast_mode.autocast('cuda', enabled=True):
        vision_outputs = clip_model(pixel_values=pixel_values, output_hidden_states=True)
        image_features = vision_outputs.hidden_states
        embedded_images = image_adapter(image_features)
        embedded_images = embedded_images.to('cuda')
    
    # Embed prompt
    prompt_embeds = text_model.model.embed_tokens(prompt.to('cuda'))
    assert prompt_embeds.shape == (1, prompt.shape[1], text_model.config.hidden_size), f"Prompt shape is {prompt_embeds.shape}, expected {(1, prompt.shape[1], text_model.config.hidden_size)}"
    embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device=text_model.device, dtype=torch.int64))
    eot_embed = image_adapter.get_eot_embedding().unsqueeze(0).to(dtype=text_model.dtype)

    # Construct prompts
    inputs_embeds = torch.cat([
        embedded_bos.expand(embedded_images.shape[0], -1, -1),
        embedded_images.to(dtype=embedded_bos.dtype),
        prompt_embeds.expand(embedded_images.shape[0], -1, -1),
        eot_embed.expand(embedded_images.shape[0], -1, -1),
    ], dim=1)

    input_ids = torch.cat([
        torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long),
        torch.zeros((1, embedded_images.shape[1]), dtype=torch.long),
        prompt,
        torch.tensor([[tokenizer.convert_tokens_to_ids("<|eot_id|>")]], dtype=torch.long),
    ], dim=1).to('cuda')
    attention_mask = torch.ones_like(input_ids)

    #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None)
    #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, suppress_tokens=None)
    generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, suppress_tokens=None)   # Uses the default which is temp=0.6, top_p=0.9

    # Trim off the prompt
    generate_ids = generate_ids[:, input_ids.shape[1]:]
    if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids("<|eot_id|>"):
        generate_ids = generate_ids[:, :-1]

    caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]

    return caption.strip()


@spaces.GPU()
@torch.no_grad()
def stream_chat_mod(input_image: Image.Image, caption_type: str, caption_tone: str, caption_length: str | int, max_new_tokens: int=300, top_p: float=0.9, temperature: float=0.6, progress=gr.Progress(track_tqdm=True)) -> str:
    global use_inference_client
    global text_model
    torch.cuda.empty_cache()
    gc.collect()

    # 'any' means no length specified
    length = None if caption_length == "any" else caption_length

    if isinstance(length, str):
        try:
            length = int(length)
        except ValueError:
            pass

    # 'rng-tags' and 'training_prompt' don't have formal/informal tones
    if caption_type == "rng-tags" or caption_type == "training_prompt":
        caption_tone = "formal"

    # Build prompt
    prompt_key = (caption_type, caption_tone, isinstance(length, str), isinstance(length, int))
    if prompt_key not in CAPTION_TYPE_MAP:
        raise ValueError(f"Invalid caption type: {prompt_key}")

    prompt_str = CAPTION_TYPE_MAP[prompt_key][0].format(length=length, word_count=length)
    print(f"Prompt: {prompt_str}")

    # Preprocess image
    #image = clip_processor(images=input_image, return_tensors='pt').pixel_values
    image = input_image.resize((384, 384), Image.LANCZOS)
    pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
    pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
    pixel_values = pixel_values.to(device)

    # Tokenize the prompt
    prompt = tokenizer.encode(prompt_str, return_tensors='pt', padding=False, truncation=False, add_special_tokens=False)

    # Embed image
    with torch.amp.autocast_mode.autocast(device, enabled=True):
        vision_outputs = clip_model(pixel_values=pixel_values, output_hidden_states=True)
        image_features = vision_outputs.hidden_states
        embedded_images = image_adapter(image_features)
        embedded_images = embedded_images.to(device)
    
    # Embed prompt
    prompt_embeds = text_model.model.embed_tokens(prompt.to(device))
    assert prompt_embeds.shape == (1, prompt.shape[1], text_model.config.hidden_size), f"Prompt shape is {prompt_embeds.shape}, expected {(1, prompt.shape[1], text_model.config.hidden_size)}"
    embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device=text_model.device, dtype=torch.int64))
    eot_embed = image_adapter.get_eot_embedding().unsqueeze(0).to(dtype=text_model.dtype)

    # Construct prompts
    inputs_embeds = torch.cat([
        embedded_bos.expand(embedded_images.shape[0], -1, -1),
        embedded_images.to(dtype=embedded_bos.dtype),
        prompt_embeds.expand(embedded_images.shape[0], -1, -1),
        eot_embed.expand(embedded_images.shape[0], -1, -1),
    ], dim=1)

    input_ids = torch.cat([
        torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long),
        torch.zeros((1, embedded_images.shape[1]), dtype=torch.long),
        prompt,
        torch.tensor([[tokenizer.convert_tokens_to_ids("<|eot_id|>")]], dtype=torch.long),
    ], dim=1).to(device)
    attention_mask = torch.ones_like(input_ids)

    text_model.to(device)
    #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None)
    #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, suppress_tokens=None)
    #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, suppress_tokens=None)   # Uses the default which is temp=0.6, top_p=0.9
    generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=max_new_tokens,
                                       do_sample=True, suppress_tokens=None, top_p=top_p, temperature=temperature)

    # Trim off the prompt
    generate_ids = generate_ids[:, input_ids.shape[1]:]
    if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids("<|eot_id|>"):
        generate_ids = generate_ids[:, :-1]

    caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]

    return caption.strip()


# https://huggingface.co/docs/transformers/v4.44.2/main_classes/text_generation#transformers.FlaxGenerationMixin.generate
# https://github.com/huggingface/transformers/issues/6535
# https://zenn.dev/hijikix/articles/8c445f4373fdcc ja
# https://github.com/ggerganov/llama.cpp/discussions/7712
# https://huggingface.co/docs/huggingface_hub/guides/inference#openai-compatibility
# https://huggingface.co/docs/huggingface_hub/v0.24.6/en/package_reference/inference_client#huggingface_hub.InferenceClient.text_generation


def is_repo_name(s):
    import re
    return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)


def is_repo_exists(repo_id):
    from huggingface_hub import HfApi
    try:
        api = HfApi(token=HF_TOKEN)
        if api.repo_exists(repo_id=repo_id): return True
        else: return False
    except Exception as e:
        print(f"Error: Failed to connect {repo_id}.")
        print(e)
        return True # for safe


def get_text_model():
    return list(llm_models.keys())


def is_gguf_repo(repo_id: str):
    from huggingface_hub import HfApi
    try:
        api = HfApi(token=HF_TOKEN)
        if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return False
        files = api.list_repo_files(repo_id=repo_id)
    except Exception as e:
        print(f"Error: Failed to get {repo_id}'s info.")
        print(e)
        gr.Warning(f"Error: Failed to get {repo_id}'s info.")
        return False
    files = [f for f in files if f.endswith(".gguf")]
    if len(files) == 0: return False
    else: return True


def get_repo_gguf(repo_id: str):
    from huggingface_hub import HfApi
    try:
        api = HfApi(token=HF_TOKEN)
        if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(value="", choices=[])
        files = api.list_repo_files(repo_id=repo_id)
    except Exception as e:
        print(f"Error: Failed to get {repo_id}'s info.")
        print(e)
        gr.Warning(f"Error: Failed to get {repo_id}'s info.")
        return gr.update(value="", choices=[])
    files = [f for f in files if f.endswith(".gguf")]
    if len(files) == 0: return gr.update(value="", choices=[])
    else: return gr.update(value=files[0], choices=files)


@spaces.GPU()
def change_text_model(model_name: str=MODEL_PATH, use_client: bool=False, gguf_file: str | None=None,

                      is_nf4: bool=True, progress=gr.Progress(track_tqdm=True)):
    global use_inference_client
    global llm_models
    use_inference_client = use_client
    try:
        if not is_repo_name(model_name) or not is_repo_exists(model_name):
            raise gr.Error(f"Repo doesn't exist: {model_name}")
        if not gguf_file and is_gguf_repo(model_name):
            gr.Info(f"Please select a gguf file.")
            return gr.update(visible=True)
        if use_inference_client:
            pass #
        else:
            load_text_model(model_name, gguf_file, is_nf4)
        if model_name not in llm_models: llm_models[model_name] = gguf_file if gguf_file else None
        return gr.update(choices=get_text_model())
    except Exception as e:
        raise gr.Error(f"Model load error: {model_name}, {e}")


# original UI
with gr.Blocks() as demo:
    gr.HTML(TITLE)

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")

            caption_type = gr.Dropdown(
                choices=["descriptive", "training_prompt", "rng-tags"],
                label="Caption Type",
                value="descriptive",
            )

            caption_tone = gr.Dropdown(
                choices=["formal", "informal"],
                label="Caption Tone",
                value="formal",
            )

            caption_length = gr.Dropdown(
                choices=["any", "very short", "short", "medium-length", "long", "very long"] +
                        [str(i) for i in range(20, 261, 10)],
                label="Caption Length",
                value="any",
            )

            gr.Markdown("**Note:** Caption tone doesn't affect `rng-tags` and `training_prompt`.")

            run_button = gr.Button("Caption")
        
        with gr.Column():
            output_caption = gr.Textbox(label="Caption")
    
    run_button.click(fn=stream_chat, inputs=[input_image, caption_type, caption_tone, caption_length], outputs=[output_caption])


if __name__ == "__main__":
    demo.launch()