import os if os.environ.get("SPACES_ZERO_GPU") is not None: import spaces else: class spaces: @staticmethod def GPU(func): def wrapper(*args, **kwargs): return func(*args, **kwargs) return wrapper import gradio as gr from huggingface_hub import InferenceClient from torch import nn from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM, LlavaForConditionalGeneration from pathlib import Path import torch import torch.amp.autocast_mode from PIL import Image import torchvision.transforms.functional as TVF import gc from peft import PeftConfig from typing import Union import subprocess subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True) BASE_DIR = Path(__file__).resolve().parent # Define the base directory device = "cuda" if torch.cuda.is_available() else "cpu" HF_TOKEN = os.environ.get("HF_TOKEN", None) use_inference_client = False PIXTRAL_PATHS = ["SeanScripts/pixtral-12b-nf4", "mistral-community/pixtral-12b"] llm_models = { "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2": None, PIXTRAL_PATHS[0]: None, "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 = BASE_DIR / Path("9em124t2-499968") LORA_PATH = CHECKPOINT_PATH / "text_model" TITLE = "

JoyCaption Alpha One (2024-09-20a)

" 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 # https://huggingface.co/docs/transformers/main/quantization/bitsandbytes?bnb=4-bit # https://huggingface.co/lllyasviel/flux1-dev-bnb-nf4 tokenizer = None text_model_client = None text_model = None image_adapter = None peft_config = None pixtral_model = None pixtral_processor = None def load_text_model(model_name: str=MODEL_PATH, gguf_file: Union[str, None]=None, is_nf4: bool=True): global tokenizer, text_model, image_adapter, peft_config, pixtral_model, pixtral_processor, text_model_client, use_inference_client try: tokenizer = None text_model_client = None text_model = None image_adapter = None peft_config = None pixtral_model = None pixtral_processor = None torch.cuda.empty_cache() gc.collect() 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) if model_name in PIXTRAL_PATHS: # Pixtral print(f"Loading LLM: {model_name}") if is_nf4: pixtral_model = LlavaForConditionalGeneration.from_pretrained(model_name, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval() else: pixtral_model = LlavaForConditionalGeneration.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval() pixtral_processor = AutoProcessor.from_pretrained(model_name) print(f"pixtral_model: {type(pixtral_model)}") # print(f"pixtral_processor: {type(pixtral_processor)}") # return 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', weights_only=True) 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(PIXTRAL_PATHS[0]) #print(f"pixtral_model: {type(pixtral_model)}") # #print(f"pixtral_processor: {type(pixtral_processor)}") # load_text_model() print(f"pixtral_model: {type(pixtral_model)}") # print(f"pixtral_processor: {type(pixtral_processor)}") # @spaces.GPU() @torch.inference_mode() def stream_chat_mod(input_image: Image.Image, caption_type: str, caption_tone: str, caption_length: Union[str, int], 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: global tokenizer, text_model, image_adapter, peft_config, pixtral_model, pixtral_processor, text_model_client, use_inference_client 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}") # Pixtral if model_name in PIXTRAL_PATHS: print(f"pixtral_model: {type(pixtral_model)}") # print(f"pixtral_processor: {type(pixtral_processor)}") # input_images = [input_image.convert("RGB")] #input_prompt = f"[INST]{prompt_str}\n[IMG][/INST]" input_prompt = "[INST]Caption this image:\n[IMG][/INST]" inputs = pixtral_processor(images=input_images, text=input_prompt, return_tensors="pt").to(device) generate_ids = pixtral_model.generate(**inputs, max_new_tokens=max_new_tokens) output = pixtral_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] return output.strip() # Preprocess image 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=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 is_valid_repo(repo_id): from huggingface_hub import HfApi import re try: if not re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', repo_id): return False api = HfApi() if api.repo_exists(repo_id=repo_id): return True else: return False except Exception as e: print(f"Failed to connect {repo_id}. {e}") return False 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: Union[str, None]=None, is_nf4: bool=True, progress=gr.Progress(track_tqdm=True)): global use_inference_client, 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}")