Upload 6 files
Browse files- README.md +14 -12
- app.py +67 -0
- joycaption.py +422 -0
- packages.txt +1 -0
- pre-requirements.txt +1 -0
- requirements.txt +12 -0
README.md
CHANGED
|
@@ -1,12 +1,14 @@
|
|
| 1 |
-
---
|
| 2 |
-
title: Joy Caption
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
-
sdk: gradio
|
| 7 |
-
sdk_version: 4.44.0
|
| 8 |
-
app_file: app.py
|
| 9 |
-
pinned: false
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: Joy Caption Alpha One Mod (dedicated to reproducing bugs, dangerous to touch)
|
| 3 |
+
emoji: 💬
|
| 4 |
+
colorFrom: yellow
|
| 5 |
+
colorTo: purple
|
| 6 |
+
sdk: gradio
|
| 7 |
+
sdk_version: 4.44.0
|
| 8 |
+
app_file: app.py
|
| 9 |
+
pinned: false
|
| 10 |
+
license: mit
|
| 11 |
+
hf_oauth: true
|
| 12 |
+
---
|
| 13 |
+
|
| 14 |
+
An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
|
app.py
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import spaces
|
| 2 |
+
import gradio as gr
|
| 3 |
+
from joycaption import stream_chat_mod, get_text_model, change_text_model, get_repo_gguf
|
| 4 |
+
|
| 5 |
+
JC_TITLE_MD = "<h1><center>JoyCaption Alpha One Mod</center></h1>"
|
| 6 |
+
JC_DESC_MD = """This space is mod of [fancyfeast/joy-caption-alpha-one](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-one),
|
| 7 |
+
[Wi-zz/joy-caption-pre-alpha](https://huggingface.co/Wi-zz/joy-caption-pre-alpha). Thanks to [dominic1021](https://huggingface.co/dominic1021)"""
|
| 8 |
+
|
| 9 |
+
css = """
|
| 10 |
+
.info {text-align:center; !important}
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
with gr.Blocks(fill_width=True, css=css, delete_cache=(60, 3600)) as demo:
|
| 14 |
+
gr.HTML(JC_TITLE_MD)
|
| 15 |
+
with gr.Row():
|
| 16 |
+
with gr.Column():
|
| 17 |
+
with gr.Group():
|
| 18 |
+
jc_input_image = gr.Image(type="pil", label="Input Image", sources=["upload", "clipboard"], height=384)
|
| 19 |
+
with gr.Row():
|
| 20 |
+
jc_caption_type = gr.Dropdown(
|
| 21 |
+
choices=["descriptive", "training_prompt", "rng-tags"],
|
| 22 |
+
label="Caption Type",
|
| 23 |
+
value="descriptive",
|
| 24 |
+
)
|
| 25 |
+
jc_caption_tone = gr.Dropdown(
|
| 26 |
+
choices=["formal", "informal"],
|
| 27 |
+
label="Caption Tone",
|
| 28 |
+
value="formal",
|
| 29 |
+
)
|
| 30 |
+
jc_caption_length = gr.Dropdown(
|
| 31 |
+
choices=["any", "very short", "short", "medium-length", "long", "very long"] +
|
| 32 |
+
[str(i) for i in range(20, 261, 10)],
|
| 33 |
+
label="Caption Length",
|
| 34 |
+
value="any",
|
| 35 |
+
)
|
| 36 |
+
gr.Markdown("**Note:** Caption tone doesn't affect `rng-tags` and `training_prompt`.", elem_classes="info")
|
| 37 |
+
with gr.Accordion("Advanced", open=False):
|
| 38 |
+
with gr.Row():
|
| 39 |
+
jc_text_model = gr.Dropdown(label="LLM Model", info="You can enter a huggingface model repo_id to want to use.",
|
| 40 |
+
choices=get_text_model(), value=get_text_model()[0],
|
| 41 |
+
allow_custom_value=True, interactive=True, min_width=320)
|
| 42 |
+
jc_gguf = gr.Dropdown(label=f"GGUF Filename", choices=[], value="",
|
| 43 |
+
allow_custom_value=True, min_width=320, visible=False)
|
| 44 |
+
jc_nf4 = gr.Checkbox(label="Use NF4 quantization", value=True)
|
| 45 |
+
jc_text_model_button = gr.Button("Load Model", variant="secondary", visible=False)
|
| 46 |
+
jc_use_inference_client = gr.Checkbox(label="Use Inference Client", value=False, visible=False)
|
| 47 |
+
with gr.Row():
|
| 48 |
+
jc_tokens = gr.Slider(minimum=1, maximum=4096, value=300, step=1, label="Max tokens")
|
| 49 |
+
jc_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.6, step=0.1, label="Temperature")
|
| 50 |
+
jc_topp = gr.Slider(minimum=0, maximum=2.0, value=0.9, step=0.01, label="Top-P")
|
| 51 |
+
jc_run_button = gr.Button("Caption", variant="primary")
|
| 52 |
+
with gr.Column():
|
| 53 |
+
jc_output_caption = gr.Textbox(label="Caption", show_copy_button=True)
|
| 54 |
+
gr.Markdown(JC_DESC_MD, elem_classes="info")
|
| 55 |
+
gr.LoginButton()
|
| 56 |
+
gr.DuplicateButton(value="Duplicate Space for private use (This demo does not work on CPU. Requires GPU Space)")
|
| 57 |
+
|
| 58 |
+
jc_run_button.click(fn=stream_chat_mod, inputs=[jc_input_image, jc_caption_type, jc_caption_tone, jc_caption_length,
|
| 59 |
+
jc_tokens, jc_topp, jc_temperature, jc_text_model], outputs=[jc_output_caption])
|
| 60 |
+
jc_text_model.change(change_text_model, [jc_text_model, jc_use_inference_client, jc_gguf, jc_nf4], [jc_text_model], show_api=False)
|
| 61 |
+
#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)
|
| 62 |
+
#jc_text_model.change(get_repo_gguf, [jc_text_model], [jc_gguf], show_api=False)
|
| 63 |
+
#jc_use_inference_client.change(change_text_model, [jc_text_model, jc_use_inference_client], [jc_text_model], show_api=False)
|
| 64 |
+
|
| 65 |
+
if __name__ == "__main__":
|
| 66 |
+
#demo.queue()
|
| 67 |
+
demo.launch()
|
joycaption.py
ADDED
|
@@ -0,0 +1,422 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
if os.environ.get("SPACES_ZERO_GPU") is not None:
|
| 3 |
+
import spaces
|
| 4 |
+
else:
|
| 5 |
+
class spaces:
|
| 6 |
+
@staticmethod
|
| 7 |
+
def GPU(func):
|
| 8 |
+
def wrapper(*args, **kwargs):
|
| 9 |
+
return func(*args, **kwargs)
|
| 10 |
+
return wrapper
|
| 11 |
+
import gradio as gr
|
| 12 |
+
from huggingface_hub import InferenceClient
|
| 13 |
+
from torch import nn
|
| 14 |
+
from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM, LlavaForConditionalGeneration
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
import torch
|
| 17 |
+
import torch.amp.autocast_mode
|
| 18 |
+
from PIL import Image
|
| 19 |
+
import torchvision.transforms.functional as TVF
|
| 20 |
+
import gc
|
| 21 |
+
from peft import PeftConfig
|
| 22 |
+
from typing import Union
|
| 23 |
+
|
| 24 |
+
import subprocess
|
| 25 |
+
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
|
| 26 |
+
|
| 27 |
+
BASE_DIR = Path(__file__).resolve().parent # Define the base directory
|
| 28 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 29 |
+
HF_TOKEN = os.environ.get("HF_TOKEN", None)
|
| 30 |
+
use_inference_client = False
|
| 31 |
+
PIXTRAL_PATHS = ["SeanScripts/pixtral-12b-nf4", "mistral-community/pixtral-12b"]
|
| 32 |
+
|
| 33 |
+
llm_models = {
|
| 34 |
+
"Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2": None,
|
| 35 |
+
PIXTRAL_PATHS[0]: None,
|
| 36 |
+
"bunnycore/LLama-3.1-8B-Matrix": None,
|
| 37 |
+
"Sao10K/Llama-3.1-8B-Stheno-v3.4": None,
|
| 38 |
+
"unsloth/Meta-Llama-3.1-8B-bnb-4bit": None,
|
| 39 |
+
"DevQuasar/HermesNova-Llama-3.1-8B": None,
|
| 40 |
+
"mergekit-community/L3.1-Boshima-b-FIX": None,
|
| 41 |
+
"meta-llama/Meta-Llama-3.1-8B": None, # gated
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
CLIP_PATH = "google/siglip-so400m-patch14-384"
|
| 45 |
+
MODEL_PATH = list(llm_models.keys())[0]
|
| 46 |
+
CHECKPOINT_PATH = BASE_DIR / Path("9em124t2-499968")
|
| 47 |
+
LORA_PATH = CHECKPOINT_PATH / "text_model"
|
| 48 |
+
TITLE = "<h1><center>JoyCaption Alpha One (2024-09-20a)</center></h1>"
|
| 49 |
+
CAPTION_TYPE_MAP = {
|
| 50 |
+
("descriptive", "formal", False, False): ["Write a descriptive caption for this image in a formal tone."],
|
| 51 |
+
("descriptive", "formal", False, True): ["Write a descriptive caption for this image in a formal tone within {word_count} words."],
|
| 52 |
+
("descriptive", "formal", True, False): ["Write a {length} descriptive caption for this image in a formal tone."],
|
| 53 |
+
("descriptive", "informal", False, False): ["Write a descriptive caption for this image in a casual tone."],
|
| 54 |
+
("descriptive", "informal", False, True): ["Write a descriptive caption for this image in a casual tone within {word_count} words."],
|
| 55 |
+
("descriptive", "informal", True, False): ["Write a {length} descriptive caption for this image in a casual tone."],
|
| 56 |
+
|
| 57 |
+
("training_prompt", "formal", False, False): ["Write a stable diffusion prompt for this image."],
|
| 58 |
+
("training_prompt", "formal", False, True): ["Write a stable diffusion prompt for this image within {word_count} words."],
|
| 59 |
+
("training_prompt", "formal", True, False): ["Write a {length} stable diffusion prompt for this image."],
|
| 60 |
+
|
| 61 |
+
("rng-tags", "formal", False, False): ["Write a list of Booru tags for this image."],
|
| 62 |
+
("rng-tags", "formal", False, True): ["Write a list of Booru tags for this image within {word_count} words."],
|
| 63 |
+
("rng-tags", "formal", True, False): ["Write a {length} list of Booru tags for this image."],
|
| 64 |
+
}
|
| 65 |
+
|
| 66 |
+
class ImageAdapter(nn.Module):
|
| 67 |
+
def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
|
| 68 |
+
super().__init__()
|
| 69 |
+
self.deep_extract = deep_extract
|
| 70 |
+
|
| 71 |
+
if self.deep_extract:
|
| 72 |
+
input_features = input_features * 5
|
| 73 |
+
|
| 74 |
+
self.linear1 = nn.Linear(input_features, output_features)
|
| 75 |
+
self.activation = nn.GELU()
|
| 76 |
+
self.linear2 = nn.Linear(output_features, output_features)
|
| 77 |
+
self.ln1 = nn.Identity() if not ln1 else nn.LayerNorm(input_features)
|
| 78 |
+
self.pos_emb = None if not pos_emb else nn.Parameter(torch.zeros(num_image_tokens, input_features))
|
| 79 |
+
|
| 80 |
+
# Mode token
|
| 81 |
+
#self.mode_token = nn.Embedding(n_modes, output_features)
|
| 82 |
+
#self.mode_token.weight.data.normal_(mean=0.0, std=0.02) # Matches HF's implementation of llama3
|
| 83 |
+
|
| 84 |
+
# Other tokens (<|image_start|>, <|image_end|>, <|eot_id|>)
|
| 85 |
+
self.other_tokens = nn.Embedding(3, output_features)
|
| 86 |
+
self.other_tokens.weight.data.normal_(mean=0.0, std=0.02) # Matches HF's implementation of llama3
|
| 87 |
+
|
| 88 |
+
def forward(self, vision_outputs: torch.Tensor):
|
| 89 |
+
if self.deep_extract:
|
| 90 |
+
x = torch.concat((
|
| 91 |
+
vision_outputs[-2],
|
| 92 |
+
vision_outputs[3],
|
| 93 |
+
vision_outputs[7],
|
| 94 |
+
vision_outputs[13],
|
| 95 |
+
vision_outputs[20],
|
| 96 |
+
), dim=-1)
|
| 97 |
+
assert len(x.shape) == 3, f"Expected 3, got {len(x.shape)}" # batch, tokens, features
|
| 98 |
+
assert x.shape[-1] == vision_outputs[-2].shape[-1] * 5, f"Expected {vision_outputs[-2].shape[-1] * 5}, got {x.shape[-1]}"
|
| 99 |
+
else:
|
| 100 |
+
x = vision_outputs[-2]
|
| 101 |
+
|
| 102 |
+
x = self.ln1(x)
|
| 103 |
+
|
| 104 |
+
if self.pos_emb is not None:
|
| 105 |
+
assert x.shape[-2:] == self.pos_emb.shape, f"Expected {self.pos_emb.shape}, got {x.shape[-2:]}"
|
| 106 |
+
x = x + self.pos_emb
|
| 107 |
+
|
| 108 |
+
x = self.linear1(x)
|
| 109 |
+
x = self.activation(x)
|
| 110 |
+
x = self.linear2(x)
|
| 111 |
+
|
| 112 |
+
# Mode token
|
| 113 |
+
#mode_token = self.mode_token(mode)
|
| 114 |
+
#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}"
|
| 115 |
+
#x = torch.cat((x, mode_token), dim=1)
|
| 116 |
+
|
| 117 |
+
# <|image_start|>, IMAGE, <|image_end|>
|
| 118 |
+
other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1))
|
| 119 |
+
assert other_tokens.shape == (x.shape[0], 2, x.shape[2]), f"Expected {(x.shape[0], 2, x.shape[2])}, got {other_tokens.shape}"
|
| 120 |
+
x = torch.cat((other_tokens[:, 0:1], x, other_tokens[:, 1:2]), dim=1)
|
| 121 |
+
|
| 122 |
+
return x
|
| 123 |
+
|
| 124 |
+
def get_eot_embedding(self):
|
| 125 |
+
return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0)
|
| 126 |
+
|
| 127 |
+
# https://huggingface.co/docs/transformers/v4.44.2/gguf
|
| 128 |
+
# https://github.com/city96/ComfyUI-GGUF/issues/7
|
| 129 |
+
# https://github.com/THUDM/ChatGLM-6B/issues/18
|
| 130 |
+
# https://github.com/meta-llama/llama/issues/394
|
| 131 |
+
# https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct/discussions/109
|
| 132 |
+
# https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu
|
| 133 |
+
# https://huggingface.co/google/flan-ul2/discussions/8
|
| 134 |
+
# https://huggingface.co/blog/4bit-transformers-bitsandbytes
|
| 135 |
+
# https://huggingface.co/docs/transformers/main/en/peft
|
| 136 |
+
# https://huggingface.co/docs/transformers/main/en/peft#enable-and-disable-adapters
|
| 137 |
+
# https://huggingface.co/docs/transformers/main/quantization/bitsandbytes?bnb=4-bit
|
| 138 |
+
# https://huggingface.co/lllyasviel/flux1-dev-bnb-nf4
|
| 139 |
+
tokenizer = None
|
| 140 |
+
text_model_client = None
|
| 141 |
+
text_model = None
|
| 142 |
+
image_adapter = None
|
| 143 |
+
peft_config = None
|
| 144 |
+
pixtral_model = None
|
| 145 |
+
pixtral_processor = None
|
| 146 |
+
def load_text_model(model_name: str=MODEL_PATH, gguf_file: Union[str, None]=None, is_nf4: bool=True):
|
| 147 |
+
global tokenizer, text_model, image_adapter, peft_config, pixtral_model, pixtral_processor, text_model_client, use_inference_client
|
| 148 |
+
try:
|
| 149 |
+
tokenizer = None
|
| 150 |
+
text_model_client = None
|
| 151 |
+
text_model = None
|
| 152 |
+
image_adapter = None
|
| 153 |
+
peft_config = None
|
| 154 |
+
pixtral_model = None
|
| 155 |
+
pixtral_processor = None
|
| 156 |
+
torch.cuda.empty_cache()
|
| 157 |
+
gc.collect()
|
| 158 |
+
|
| 159 |
+
from transformers import BitsAndBytesConfig
|
| 160 |
+
nf4_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
|
| 161 |
+
bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
|
| 162 |
+
|
| 163 |
+
if model_name in PIXTRAL_PATHS: # Pixtral
|
| 164 |
+
print(f"Loading LLM: {model_name}")
|
| 165 |
+
if is_nf4:
|
| 166 |
+
pixtral_model = LlavaForConditionalGeneration.from_pretrained(model_name, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 167 |
+
else:
|
| 168 |
+
pixtral_model = LlavaForConditionalGeneration.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 169 |
+
pixtral_processor = AutoProcessor.from_pretrained(model_name)
|
| 170 |
+
print(f"pixtral_model: {type(pixtral_model)}") #
|
| 171 |
+
print(f"pixtral_processor: {type(pixtral_processor)}") #
|
| 172 |
+
return
|
| 173 |
+
|
| 174 |
+
print("Loading tokenizer")
|
| 175 |
+
if gguf_file: tokenizer = AutoTokenizer.from_pretrained(model_name, gguf_file=gguf_file, use_fast=True, legacy=False)
|
| 176 |
+
else: tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, legacy=False)
|
| 177 |
+
assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
|
| 178 |
+
|
| 179 |
+
print(f"Loading LLM: {model_name}")
|
| 180 |
+
if gguf_file:
|
| 181 |
+
if device == "cpu":
|
| 182 |
+
text_model = AutoModelForCausalLM.from_pretrained(model_name, gguf_file=gguf_file, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 183 |
+
elif is_nf4:
|
| 184 |
+
text_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 185 |
+
else:
|
| 186 |
+
text_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 187 |
+
else:
|
| 188 |
+
if device == "cpu":
|
| 189 |
+
text_model = AutoModelForCausalLM.from_pretrained(model_name, gguf_file=gguf_file, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 190 |
+
elif is_nf4:
|
| 191 |
+
text_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 192 |
+
else:
|
| 193 |
+
text_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval()
|
| 194 |
+
|
| 195 |
+
if LORA_PATH.exists():
|
| 196 |
+
print("Loading VLM's custom text model")
|
| 197 |
+
if is_nf4: peft_config = PeftConfig.from_pretrained(LORA_PATH, device_map=device, quantization_config=nf4_config)
|
| 198 |
+
else: peft_config = PeftConfig.from_pretrained(LORA_PATH, device_map=device)
|
| 199 |
+
text_model.add_adapter(peft_config)
|
| 200 |
+
text_model.enable_adapters()
|
| 201 |
+
|
| 202 |
+
print("Loading image adapter")
|
| 203 |
+
image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False).eval().to("cpu")
|
| 204 |
+
image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu", weights_only=True))
|
| 205 |
+
image_adapter.eval().to(device)
|
| 206 |
+
except Exception as e:
|
| 207 |
+
print(f"LLM load error: {e}")
|
| 208 |
+
raise Exception(f"LLM load error: {e}") from e
|
| 209 |
+
finally:
|
| 210 |
+
torch.cuda.empty_cache()
|
| 211 |
+
gc.collect()
|
| 212 |
+
|
| 213 |
+
load_text_model.zerogpu = True
|
| 214 |
+
|
| 215 |
+
# Load CLIP
|
| 216 |
+
print("Loading CLIP")
|
| 217 |
+
clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
|
| 218 |
+
clip_model = AutoModel.from_pretrained(CLIP_PATH).vision_model
|
| 219 |
+
if (CHECKPOINT_PATH / "clip_model.pt").exists():
|
| 220 |
+
print("Loading VLM's custom vision model")
|
| 221 |
+
checkpoint = torch.load(CHECKPOINT_PATH / "clip_model.pt", map_location='cpu', weights_only=True)
|
| 222 |
+
checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()}
|
| 223 |
+
clip_model.load_state_dict(checkpoint)
|
| 224 |
+
del checkpoint
|
| 225 |
+
clip_model.eval().requires_grad_(False).to(device)
|
| 226 |
+
|
| 227 |
+
# Tokenizer
|
| 228 |
+
# LLM
|
| 229 |
+
# Image Adapter
|
| 230 |
+
#load_text_model(PIXTRAL_PATHS[0])
|
| 231 |
+
#print(f"pixtral_model: {type(pixtral_model)}") #
|
| 232 |
+
#print(f"pixtral_processor: {type(pixtral_processor)}") #
|
| 233 |
+
load_text_model()
|
| 234 |
+
print(f"pixtral_model: {type(pixtral_model)}") #
|
| 235 |
+
print(f"pixtral_processor: {type(pixtral_processor)}") #
|
| 236 |
+
|
| 237 |
+
@spaces.GPU()
|
| 238 |
+
@torch.inference_mode()
|
| 239 |
+
def stream_chat_mod(input_image: Image.Image, caption_type: str, caption_tone: str, caption_length: Union[str, int],
|
| 240 |
+
max_new_tokens: int=300, top_p: float=0.9, temperature: float=0.6, model_name: str=MODEL_PATH, progress=gr.Progress(track_tqdm=True)) -> str:
|
| 241 |
+
global tokenizer, text_model, image_adapter, peft_config, pixtral_model, pixtral_processor, text_model_client, use_inference_client
|
| 242 |
+
torch.cuda.empty_cache()
|
| 243 |
+
gc.collect()
|
| 244 |
+
|
| 245 |
+
# 'any' means no length specified
|
| 246 |
+
length = None if caption_length == "any" else caption_length
|
| 247 |
+
|
| 248 |
+
if isinstance(length, str):
|
| 249 |
+
try:
|
| 250 |
+
length = int(length)
|
| 251 |
+
except ValueError:
|
| 252 |
+
pass
|
| 253 |
+
|
| 254 |
+
# 'rng-tags' and 'training_prompt' don't have formal/informal tones
|
| 255 |
+
if caption_type == "rng-tags" or caption_type == "training_prompt":
|
| 256 |
+
caption_tone = "formal"
|
| 257 |
+
|
| 258 |
+
# Build prompt
|
| 259 |
+
prompt_key = (caption_type, caption_tone, isinstance(length, str), isinstance(length, int))
|
| 260 |
+
if prompt_key not in CAPTION_TYPE_MAP:
|
| 261 |
+
raise ValueError(f"Invalid caption type: {prompt_key}")
|
| 262 |
+
|
| 263 |
+
prompt_str = CAPTION_TYPE_MAP[prompt_key][0].format(length=length, word_count=length)
|
| 264 |
+
print(f"Prompt: {prompt_str}")
|
| 265 |
+
|
| 266 |
+
# Pixtral
|
| 267 |
+
if model_name in PIXTRAL_PATHS:
|
| 268 |
+
print(f"pixtral_model: {type(pixtral_model)}") #
|
| 269 |
+
print(f"pixtral_processor: {type(pixtral_processor)}") #
|
| 270 |
+
input_images = [input_image.convert("RGB")]
|
| 271 |
+
#input_prompt = f"[INST]{prompt_str}\n[IMG][/INST]"
|
| 272 |
+
input_prompt = "[INST]Caption this image:\n[IMG][/INST]"
|
| 273 |
+
inputs = pixtral_processor(images=input_images, text=input_prompt, return_tensors="pt").to(device)
|
| 274 |
+
generate_ids = pixtral_model.generate(**inputs, max_new_tokens=max_new_tokens)
|
| 275 |
+
output = pixtral_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
| 276 |
+
return output.strip()
|
| 277 |
+
|
| 278 |
+
# Preprocess image
|
| 279 |
+
image = input_image.resize((384, 384), Image.LANCZOS)
|
| 280 |
+
pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
|
| 281 |
+
pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
|
| 282 |
+
pixel_values = pixel_values.to(device)
|
| 283 |
+
|
| 284 |
+
# Tokenize the prompt
|
| 285 |
+
prompt = tokenizer.encode(prompt_str, return_tensors='pt', padding=False, truncation=False, add_special_tokens=False)
|
| 286 |
+
|
| 287 |
+
# Embed image
|
| 288 |
+
with torch.amp.autocast_mode.autocast(device, enabled=True):
|
| 289 |
+
vision_outputs = clip_model(pixel_values=pixel_values, output_hidden_states=True)
|
| 290 |
+
image_features = vision_outputs.hidden_states
|
| 291 |
+
embedded_images = image_adapter(image_features)
|
| 292 |
+
embedded_images = embedded_images.to(device)
|
| 293 |
+
|
| 294 |
+
# Embed prompt
|
| 295 |
+
prompt_embeds = text_model.model.embed_tokens(prompt.to(device))
|
| 296 |
+
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)}"
|
| 297 |
+
embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device=text_model.device, dtype=torch.int64))
|
| 298 |
+
eot_embed = image_adapter.get_eot_embedding().unsqueeze(0).to(dtype=text_model.dtype)
|
| 299 |
+
|
| 300 |
+
# Construct prompts
|
| 301 |
+
inputs_embeds = torch.cat([
|
| 302 |
+
embedded_bos.expand(embedded_images.shape[0], -1, -1),
|
| 303 |
+
embedded_images.to(dtype=embedded_bos.dtype),
|
| 304 |
+
prompt_embeds.expand(embedded_images.shape[0], -1, -1),
|
| 305 |
+
eot_embed.expand(embedded_images.shape[0], -1, -1),
|
| 306 |
+
], dim=1)
|
| 307 |
+
|
| 308 |
+
input_ids = torch.cat([
|
| 309 |
+
torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long),
|
| 310 |
+
torch.zeros((1, embedded_images.shape[1]), dtype=torch.long),
|
| 311 |
+
prompt,
|
| 312 |
+
torch.tensor([[tokenizer.convert_tokens_to_ids("<|eot_id|>")]], dtype=torch.long),
|
| 313 |
+
], dim=1).to(device)
|
| 314 |
+
attention_mask = torch.ones_like(input_ids)
|
| 315 |
+
|
| 316 |
+
text_model.to(device)
|
| 317 |
+
generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=max_new_tokens,
|
| 318 |
+
do_sample=True, suppress_tokens=None, top_p=top_p, temperature=temperature)
|
| 319 |
+
|
| 320 |
+
# Trim off the prompt
|
| 321 |
+
generate_ids = generate_ids[:, input_ids.shape[1]:]
|
| 322 |
+
if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids("<|eot_id|>"):
|
| 323 |
+
generate_ids = generate_ids[:, :-1]
|
| 324 |
+
|
| 325 |
+
caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]
|
| 326 |
+
|
| 327 |
+
return caption.strip()
|
| 328 |
+
|
| 329 |
+
|
| 330 |
+
# https://huggingface.co/docs/transformers/v4.44.2/main_classes/text_generation#transformers.FlaxGenerationMixin.generate
|
| 331 |
+
# https://github.com/huggingface/transformers/issues/6535
|
| 332 |
+
# https://zenn.dev/hijikix/articles/8c445f4373fdcc ja
|
| 333 |
+
# https://github.com/ggerganov/llama.cpp/discussions/7712
|
| 334 |
+
# https://huggingface.co/docs/huggingface_hub/guides/inference#openai-compatibility
|
| 335 |
+
# https://huggingface.co/docs/huggingface_hub/v0.24.6/en/package_reference/inference_client#huggingface_hub.InferenceClient.text_generation
|
| 336 |
+
|
| 337 |
+
|
| 338 |
+
def is_repo_name(s):
|
| 339 |
+
import re
|
| 340 |
+
return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
|
| 341 |
+
|
| 342 |
+
|
| 343 |
+
def is_repo_exists(repo_id):
|
| 344 |
+
from huggingface_hub import HfApi
|
| 345 |
+
try:
|
| 346 |
+
api = HfApi(token=HF_TOKEN)
|
| 347 |
+
if api.repo_exists(repo_id=repo_id): return True
|
| 348 |
+
else: return False
|
| 349 |
+
except Exception as e:
|
| 350 |
+
print(f"Error: Failed to connect {repo_id}.")
|
| 351 |
+
print(e)
|
| 352 |
+
return True # for safe
|
| 353 |
+
|
| 354 |
+
|
| 355 |
+
def is_valid_repo(repo_id):
|
| 356 |
+
from huggingface_hub import HfApi
|
| 357 |
+
import re
|
| 358 |
+
try:
|
| 359 |
+
if not re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', repo_id): return False
|
| 360 |
+
api = HfApi()
|
| 361 |
+
if api.repo_exists(repo_id=repo_id): return True
|
| 362 |
+
else: return False
|
| 363 |
+
except Exception as e:
|
| 364 |
+
print(f"Failed to connect {repo_id}. {e}")
|
| 365 |
+
return False
|
| 366 |
+
|
| 367 |
+
|
| 368 |
+
def get_text_model():
|
| 369 |
+
return list(llm_models.keys())
|
| 370 |
+
|
| 371 |
+
|
| 372 |
+
def is_gguf_repo(repo_id: str):
|
| 373 |
+
from huggingface_hub import HfApi
|
| 374 |
+
try:
|
| 375 |
+
api = HfApi(token=HF_TOKEN)
|
| 376 |
+
if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return False
|
| 377 |
+
files = api.list_repo_files(repo_id=repo_id)
|
| 378 |
+
except Exception as e:
|
| 379 |
+
print(f"Error: Failed to get {repo_id}'s info.")
|
| 380 |
+
print(e)
|
| 381 |
+
gr.Warning(f"Error: Failed to get {repo_id}'s info.")
|
| 382 |
+
return False
|
| 383 |
+
files = [f for f in files if f.endswith(".gguf")]
|
| 384 |
+
if len(files) == 0: return False
|
| 385 |
+
else: return True
|
| 386 |
+
|
| 387 |
+
|
| 388 |
+
def get_repo_gguf(repo_id: str):
|
| 389 |
+
from huggingface_hub import HfApi
|
| 390 |
+
try:
|
| 391 |
+
api = HfApi(token=HF_TOKEN)
|
| 392 |
+
if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(value="", choices=[])
|
| 393 |
+
files = api.list_repo_files(repo_id=repo_id)
|
| 394 |
+
except Exception as e:
|
| 395 |
+
print(f"Error: Failed to get {repo_id}'s info.")
|
| 396 |
+
print(e)
|
| 397 |
+
gr.Warning(f"Error: Failed to get {repo_id}'s info.")
|
| 398 |
+
return gr.update(value="", choices=[])
|
| 399 |
+
files = [f for f in files if f.endswith(".gguf")]
|
| 400 |
+
if len(files) == 0: return gr.update(value="", choices=[])
|
| 401 |
+
else: return gr.update(value=files[0], choices=files)
|
| 402 |
+
|
| 403 |
+
|
| 404 |
+
@spaces.GPU()
|
| 405 |
+
def change_text_model(model_name: str=MODEL_PATH, use_client: bool=False, gguf_file: Union[str, None]=None,
|
| 406 |
+
is_nf4: bool=True, progress=gr.Progress(track_tqdm=True)):
|
| 407 |
+
global use_inference_client, llm_models
|
| 408 |
+
use_inference_client = use_client
|
| 409 |
+
try:
|
| 410 |
+
if not is_repo_name(model_name) or not is_repo_exists(model_name):
|
| 411 |
+
raise gr.Error(f"Repo doesn't exist: {model_name}")
|
| 412 |
+
if not gguf_file and is_gguf_repo(model_name):
|
| 413 |
+
gr.Info(f"Please select a gguf file.")
|
| 414 |
+
return gr.update(visible=True)
|
| 415 |
+
if use_inference_client:
|
| 416 |
+
pass #
|
| 417 |
+
else:
|
| 418 |
+
load_text_model(model_name, gguf_file, is_nf4)
|
| 419 |
+
if model_name not in llm_models: llm_models[model_name] = gguf_file if gguf_file else None
|
| 420 |
+
return gr.update(choices=get_text_model())
|
| 421 |
+
except Exception as e:
|
| 422 |
+
raise gr.Error(f"Model load error: {model_name}, {e}")
|
packages.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
git-lfs
|
pre-requirements.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
pip>=23.0.0
|
requirements.txt
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
huggingface_hub
|
| 2 |
+
accelerate
|
| 3 |
+
torch
|
| 4 |
+
git+https://github.com/huggingface/transformers
|
| 5 |
+
sentencepiece
|
| 6 |
+
bitsandbytes
|
| 7 |
+
Pillow
|
| 8 |
+
protobuf
|
| 9 |
+
gguf
|
| 10 |
+
numpy<2.0.0
|
| 11 |
+
peft
|
| 12 |
+
torchvision
|