Upload 12 files

9em124t2-499968/clip_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d7b0548d12fa649370896982c2af9d03d43285b782bd47639c96e6e0b29473c
+size 1713067838

9em124t2-499968/config.yaml

@@ -0,0 +1,39 @@
+wandb_project: joy-caption-1
+device_batch_size: 2
+batch_size: 256
+learning_rate: 0.0002
+warmup_samples: 18000
+max_samples: 500000
+save_every: 50000
+test_every: 50000
+use_amp: true
+grad_scaler: true
+lr_scheduler_type: cosine
+min_lr_ratio: 0.0
+allow_tf32: true
+seed: 69
+num_workers: 8
+optimizer_type: adamw
+adam_beta1: 0.9
+adam_beta2: 0.999
+adam_eps: 1.0e-08
+adam_weight_decay: 0.0
+clip_grad_norm: 1.0
+dataset: fancyfeast/joy-captioning-20240917a
+clip_model: google/siglip-so400m-patch14-384
+text_model: meta-llama/Meta-Llama-3.1-8B
+resume: null
+gradient_checkpointing: false
+test_size: 2048
+grad_scaler_init: 65536.0
+max_caption_length: 257
+num_image_tokens: 32
+adapter_type: mlp
+text_model_dtype: bfloat16
+pre_test: false
+train_image_model: true
+image_model_lr: null
+train_lora: true
+lora_r: 64
+lora_alpha: 16
+lora_dropout: 0.1

9em124t2-499968/image_adapter.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e53c3bf8df745a3c19ae3c70dbf9bf23cfdc8f3fdb937000a4eafd2a36914661
+size 86067714

9em124t2-499968/text_model/README.md

@@ -0,0 +1,202 @@
+---
+base_model: meta-llama/Meta-Llama-3.1-8B
+library_name: peft
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.12.0

9em124t2-499968/text_model/adapter_config.json

@@ -0,0 +1,29 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "meta-llama/Meta-Llama-3.1-8B",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "q_proj",
+    "v_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

9em124t2-499968/text_model/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b48221de174ab0db7b46b4833118c5c0a4c2bf0b51b77b4cc4ab04651bd06cca
+size 109069176

README.md

@@ -1,10 +1,10 @@
 ---
-title: Joy Caption Pre Alpha Mod
+title: Joy Caption Alpha One Mod
 emoji: 💬
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.43.0
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -2,20 +2,38 @@ import spaces
 import gradio as gr
 from joycaption import stream_chat_mod, get_text_model, change_text_model, get_repo_gguf
 
-JC_TITLE_MD = "<h1><center>JoyCaption Pre-Alpha Mod</center></h1>"
-JC_DESC_MD = """This space is mod of [fancyfeast/joy-caption-pre-alpha](https://huggingface.co/spaces/fancyfeast/joy-caption-pre-alpha),
+JC_TITLE_MD = "<h1><center>JoyCaption Alpha One Mod</center></h1>"
+JC_DESC_MD = """This space is mod of [fancyfeast/joy-caption-alpha-one](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-one),
  [Wi-zz/joy-caption-pre-alpha](https://huggingface.co/Wi-zz/joy-caption-pre-alpha)"""
 
 css = """
-.info {text-align:center; display:inline-flex; align-items:center !important}
+.info {text-align:center; !important}
 """
 
-with gr.Blocks(delete_cache=(60, 3600)) as demo:
+with gr.Blocks(fill_width=True, css=css, delete_cache=(60, 3600)) as demo:
     gr.HTML(JC_TITLE_MD)
     with gr.Row():
         with gr.Column():
             with gr.Group():
                 jc_input_image = gr.Image(type="pil", label="Input Image", sources=["upload", "clipboard"], height=384)
+                with gr.Row():
+                    jc_caption_type = gr.Dropdown(
+                        choices=["descriptive", "training_prompt", "rng-tags"],
+                        label="Caption Type",
+                        value="descriptive",
+                    )
+                    jc_caption_tone = gr.Dropdown(
+                        choices=["formal", "informal"],
+                        label="Caption Tone",
+                        value="formal",
+                    )
+                    jc_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`.", elem_classes="info")
                 with gr.Accordion("Advanced", open=False):
                     with gr.Row():
                         jc_text_model = gr.Dropdown(label="LLM Model", info="You can enter a huggingface model repo_id to want to use.",
@@ -28,8 +46,8 @@ with gr.Blocks(delete_cache=(60, 3600)) as demo:
                     jc_use_inference_client = gr.Checkbox(label="Use Inference Client", value=False, visible=False)
                     with gr.Row():
                         jc_tokens = gr.Slider(minimum=1, maximum=4096, value=300, step=1, label="Max tokens")
-                        jc_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.5, step=0.1, label="Temperature")
-                        jc_topk = gr.Slider(minimum=0, maximum=100, value=40, step=10, label="Top-k")
+                        jc_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.6, step=0.1, label="Temperature")
+                        jc_topp = gr.Slider(minimum=0, maximum=2.0, value=0.9, step=0.01, label="Top-P")
             jc_run_button = gr.Button("Caption", variant="primary")
 
         with gr.Column():
@@ -38,11 +56,11 @@ with gr.Blocks(delete_cache=(60, 3600)) as demo:
     gr.LoginButton()
     gr.DuplicateButton(value="Duplicate Space for private use (This demo does not work on CPU. Requires GPU Space)")
 
-    jc_run_button.click(fn=stream_chat_mod, inputs=[jc_input_image, jc_tokens, jc_topk, jc_temperature], outputs=[jc_output_caption])
+    jc_run_button.click(fn=stream_chat_mod, inputs=[jc_input_image, jc_caption_type, jc_caption_tone, jc_caption_length, jc_tokens, jc_topp, jc_temperature], outputs=[jc_output_caption])
     jc_text_model_button.click(change_text_model, [jc_text_model, jc_use_inference_client, jc_gguf, jc_nf4], [jc_text_model], show_api=False)
     #jc_text_model.change(get_repo_gguf, [jc_text_model], [jc_gguf], show_api=False)
     jc_use_inference_client.change(change_text_model, [jc_text_model, jc_use_inference_client], [jc_text_model], show_api=False)
 
 if __name__ == "__main__":
-    demo.queue()
+    #demo.queue()
     demo.launch()

joycaption.py

@@ -8,7 +8,9 @@ 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)
@@ -24,24 +26,89 @@ llm_models = {
 }
 
 CLIP_PATH = "google/siglip-so400m-patch14-384"
-VLM_PROMPT = "A descriptive caption for this image:\n"
 MODEL_PATH = list(llm_models.keys())[0]
-CHECKPOINT_PATH = Path("wpkklhc6")
-TITLE = "<h1><center>JoyCaption Pre-Alpha (2024-07-30a)</center></h1>"
+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):
+    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):
-        x = self.linear1(vision_outputs)
+        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
@@ -50,14 +117,18 @@ class ImageAdapter(nn.Module):
 # 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:
@@ -77,8 +148,14 @@ def load_text_model(model_name: str=MODEL_PATH, gguf_file: str | None=None, is_n
             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).eval().to("cpu")
+        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:
@@ -93,57 +170,95 @@ 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.eval().requires_grad_(False).to(device)
+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):
+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 = image.to(device)
+    #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(VLM_PROMPT, return_tensors='pt', padding=False, truncation=False, add_special_tokens=False)
+    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=image, output_hidden_states=True)
-        image_features = vision_outputs.hidden_states[-2]
+    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(device)
+        embedded_images = embedded_images.to('cuda')
     
     # Embed prompt
-    prompt_embeds = text_model.model.embed_tokens(prompt.to(device))
+    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,
-    ], dim=1).to(device)
+        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, 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:
+    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]
@@ -153,23 +268,47 @@ def stream_chat(input_image: Image.Image):
 
 @spaces.GPU()
 @torch.no_grad()
-def stream_chat_mod(input_image: Image.Image, max_new_tokens: int=300, top_k: int=10, temperature: float=0.5, progress=gr.Progress(track_tqdm=True)):
+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 = image.to(device)
+    #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(VLM_PROMPT, return_tensors='pt', padding=False, truncation=False, add_special_tokens=False)
+    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=image, output_hidden_states=True)
-        image_features = vision_outputs.hidden_states[-2]
+        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)
     
@@ -177,34 +316,34 @@ def stream_chat_mod(input_image: Image.Image, max_new_tokens: int=300, top_k: in
     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)
 
-    # 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
+    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=max_new_tokens, do_sample=True, top_k=top_k, temperature=temperature, 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:
+    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]
@@ -212,6 +351,14 @@ def stream_chat_mod(input_image: Image.Image, max_new_tokens: int=300, top_k: in
     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)
@@ -290,16 +437,39 @@ def change_text_model(model_name: str=MODEL_PATH, use_client: bool=False, gguf_f
 # 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], outputs=[output_caption])
+    run_button.click(fn=stream_chat, inputs=[input_image, caption_type, caption_tone, caption_length], outputs=[output_caption])
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

requirements.txt

@@ -7,4 +7,6 @@ bitsandbytes
 Pillow
 protobuf
 gguf
-numpy<2.0.0
+numpy<2.0.0
+peft
+torchvision