Upload 4 files

README.md

@@ -1,13 +1,13 @@
----
-title: Joy Caption Pre Alpha
-emoji: 💬
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 4.36.1
-app_file: app.py
-pinned: false
-license: mit
----
-
+---
+title: Joy Caption Pre Alpha Mod
+emoji: 💬
+colorFrom: yellow
+colorTo: purple
+sdk: gradio
+sdk_version: 4.43.0
+app_file: app.py
+pinned: false
+license: mit
+---
+
 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

@@ -1,129 +1,40 @@
-import spaces
-import gradio as gr
-from huggingface_hub import InferenceClient
-from torch import nn
-from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
-from pathlib import Path
-import torch
-import torch.amp.autocast_mode
-from PIL import Image
-import os
-
-
-CLIP_PATH = "google/siglip-so400m-patch14-384"
-VLM_PROMPT = "A descriptive caption for this image:\n"
-MODEL_PATH = "meta-llama/Meta-Llama-3.1-8B"
-CHECKPOINT_PATH = Path("wpkklhc6")
-TITLE = "<h1><center>JoyCaption Pre-Alpha (2024-07-30a)</center></h1>"
-
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-
-
-class ImageAdapter(nn.Module):
-	def __init__(self, input_features: int, output_features: int):
-		super().__init__()
-		self.linear1 = nn.Linear(input_features, output_features)
-		self.activation = nn.GELU()
-		self.linear2 = nn.Linear(output_features, output_features)
-	
-	def forward(self, vision_outputs: torch.Tensor):
-		x = self.linear1(vision_outputs)
-		x = self.activation(x)
-		x = self.linear2(x)
-		return x
-
-
-# Load CLIP
-print("Loading CLIP")
-clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
-clip_model = AutoModel.from_pretrained(CLIP_PATH)
-clip_model = clip_model.vision_model
-clip_model.eval()
-clip_model.requires_grad_(False)
-clip_model.to("cuda")
-
-
-# Tokenizer
-print("Loading tokenizer")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False)
-assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
-
-# LLM
-print("Loading LLM")
-text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.bfloat16)
-text_model.eval()
-
-# Image Adapter
-print("Loading image adapter")
-image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size)
-image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu"))
-image_adapter.eval()
-image_adapter.to("cuda")
-
-
-@spaces.GPU()
-@torch.no_grad()
-def stream_chat(input_image: Image.Image):
-	torch.cuda.empty_cache()
-
-	# Preprocess image
-	image = clip_processor(images=input_image, return_tensors='pt').pixel_values
-	image = image.to('cuda')
-
-	# Tokenize the prompt
-	prompt = tokenizer.encode(VLM_PROMPT, return_tensors='pt', padding=False, truncation=False, add_special_tokens=False)
-
-	# Embed image
-	with torch.amp.autocast_mode.autocast('cuda', enabled=True):
-		vision_outputs = clip_model(pixel_values=image, output_hidden_states=True)
-		image_features = vision_outputs.hidden_states[-2]
-		embedded_images = image_adapter(image_features)
-		embedded_images = embedded_images.to('cuda')
-	
-	# Embed prompt
-	prompt_embeds = text_model.model.embed_tokens(prompt.to('cuda'))
-	assert prompt_embeds.shape == (1, prompt.shape[1], text_model.config.hidden_size), f"Prompt shape is {prompt_embeds.shape}, expected {(1, prompt.shape[1], text_model.config.hidden_size)}"
-	embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device=text_model.device, dtype=torch.int64))
-
-	# 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),
-	], 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('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)
-
-	# Trim off the prompt
-	generate_ids = generate_ids[:, input_ids.shape[1]:]
-	if generate_ids[0][-1] == tokenizer.eos_token_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()
-
-
-with gr.Blocks() as demo:
-	gr.HTML(TITLE)
-	with gr.Row():
-		with gr.Column():
-			input_image = gr.Image(type="pil", label="Input Image")
-			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])
-
-
-if __name__ == "__main__":
-    demo.launch()
+import spaces
+import gradio as gr
+from joycaption import stream_chat_mod, get_text_model, change_text_model
+
+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),
+ [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}
+"""
+
+with gr.Blocks() 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.Accordion("Advanced", open=False):
+                    jc_text_model = gr.Dropdown(label="LLM Model", info="You can enter a huggingface model repo_id to want to use.",
+                                                choices=get_text_model(), value=get_text_model()[0],
+                                                allow_custom_value=True, interactive=True, min_width=320)
+                    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_run_button = gr.Button("Caption", variant="primary")
+
+        with gr.Column():
+            jc_output_caption = gr.Textbox(label="Caption", show_copy_button=True)
+    gr.Markdown(JC_DESC_MD, elem_classes="info")
+
+    jc_run_button.click(fn=stream_chat_mod, inputs=[jc_input_image, jc_tokens, jc_topk, jc_temperature], outputs=[jc_output_caption])
+    jc_text_model.change(change_text_model, [jc_text_model, jc_use_inference_client], [jc_text_model], 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.launch()

joycaption.py

@@ -0,0 +1,250 @@
+import spaces
+import gradio as gr
+from huggingface_hub import InferenceClient
+from torch import nn
+from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
+from pathlib import Path
+import torch
+import torch.amp.autocast_mode
+from PIL import Image
+import os
+import gc
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+llm_models = [
+    "Sao10K/Llama-3.1-8B-Stheno-v3.4",
+    "unsloth/Meta-Llama-3.1-8B-bnb-4bit",
+    "mergekit-community/L3.1-Boshima-b-FIX",
+    "meta-llama/Meta-Llama-3.1-8B",
+]
+
+
+CLIP_PATH = "google/siglip-so400m-patch14-384"
+VLM_PROMPT = "A descriptive caption for this image:\n"
+MODEL_PATH = llm_models[0]
+CHECKPOINT_PATH = Path("wpkklhc6")
+TITLE = "<h1><center>JoyCaption Pre-Alpha (2024-07-30a)</center></h1>"
+
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+use_inference_client = False
+
+class ImageAdapter(nn.Module):
+    def __init__(self, input_features: int, output_features: int):
+        super().__init__()
+        self.linear1 = nn.Linear(input_features, output_features)
+        self.activation = nn.GELU()
+        self.linear2 = nn.Linear(output_features, output_features)
+    
+    def forward(self, vision_outputs: torch.Tensor):
+        x = self.linear1(vision_outputs)
+        x = self.activation(x)
+        x = self.linear2(x)
+        return x
+
+# https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu
+# https://huggingface.co/google/flan-ul2/discussions/8
+
+text_model_client = None
+text_model = None
+image_adapter = None
+def load_text_model(model_name: str=MODEL_PATH):
+    global text_model
+    global image_adapter
+    global text_model_client
+    global use_inference_client
+    try:
+        print(f"Loading LLM: {model_name}")
+        if device == "cpu": text_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval()
+        else: text_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16).eval()
+        print("Loading image adapter")
+        image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size).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.eval().requires_grad_(False).to(device)
+
+# Tokenizer
+print("Loading tokenizer")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False)
+assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
+
+# LLM
+# Image Adapter
+load_text_model()
+
+@spaces.GPU()
+@torch.no_grad()
+def stream_chat(input_image: Image.Image):
+    torch.cuda.empty_cache()
+
+    # Preprocess image
+    image = clip_processor(images=input_image, return_tensors='pt').pixel_values
+    image = image.to(device)
+
+    # Tokenize the prompt
+    prompt = tokenizer.encode(VLM_PROMPT, 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]
+        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))
+
+    # 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),
+    ], 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)
+    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)
+
+    # Trim off the prompt
+    generate_ids = generate_ids[:, input_ids.shape[1]:]
+    if generate_ids[0][-1] == tokenizer.eos_token_id:
+        generate_ids = generate_ids[:, :-1]
+
+    caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]
+
+    return caption.strip()
+
+
+@spaces.GPU()
+@torch.no_grad()
+def stream_chat_mod(input_image: Image.Image, max_new_tokens: int=300, top_k: int=10, temperature: float=0.5, progress=gr.Progress(track_tqdm=True)):
+    global use_inference_client
+    global text_model
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    # Preprocess image
+    image = clip_processor(images=input_image, return_tensors='pt').pixel_values
+    image = image.to(device)
+
+    # Tokenize the prompt
+    prompt = tokenizer.encode(VLM_PROMPT, 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]
+        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))
+
+    # 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),
+    ], 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)
+    attention_mask = torch.ones_like(input_ids)
+
+    # 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
+    #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)
+
+    # Trim off the prompt
+    generate_ids = generate_ids[:, input_ids.shape[1]:]
+    if generate_ids[0][-1] == tokenizer.eos_token_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()
+
+
+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
+    api = HfApi()
+    try:
+        if api.repo_exists(repo_id=repo_id): return True
+        else: return False
+    except Exception as e:
+        print(f"Error: Failed to connect {repo_id}.")
+        print(e)
+        return True # for safe
+
+
+def get_text_model():
+    return llm_models
+
+
+@spaces.GPU()
+def change_text_model(model_name: str=MODEL_PATH, use_client: bool=False, progress=gr.Progress(track_tqdm=True)):
+    global use_inference_client
+    global text_model
+    global llm_models
+    use_inference_client = use_client
+    try:
+        if not is_repo_name(model_name) or not is_repo_exists(model_name):
+            raise gr.Error(f"Repo doesn't exist: {model_name}")
+        if use_inference_client:
+            pass
+        else:
+            load_text_model(model_name)
+        if model_name not in llm_models: llm_models.append(model_name)
+        return gr.update(visible=True)
+    except Exception as e:
+        raise gr.Error(f"Model load error: {model_name}, {e}")
+
+
+# original UI
+with gr.Blocks() as demo:
+    gr.HTML(TITLE)
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            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])
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,5 +1,8 @@
-huggingface_hub==0.24.3
-accelerate
-torch
-transformers==4.43.3
-sentencepiece
+huggingface_hub
+accelerate
+torch
+transformers==4.43.3
+sentencepiece
+bitsandbytes
+Pillow
+protobuf