John6666 commited on
Commit
3c411d3
β€’
1 Parent(s): 01ed83d

Upload 6 files

Browse files
Files changed (6) hide show
  1. README.md +14 -12
  2. app.py +67 -0
  3. joycaption.py +422 -0
  4. packages.txt +1 -0
  5. pre-requirements.txt +1 -0
  6. requirements.txt +12 -0
README.md CHANGED
@@ -1,12 +1,14 @@
1
- ---
2
- title: Joy Caption Pre Alpha Mod Error
3
- emoji: πŸ’»
4
- colorFrom: blue
5
- colorTo: yellow
6
- sdk: gradio
7
- sdk_version: 4.44.0
8
- app_file: app.py
9
- pinned: false
10
- ---
11
-
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
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