Initial commit

app.py

@@ -0,0 +1,199 @@
+import spaces
+import gradio as gr
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, LlavaForConditionalGeneration, TextIteratorStreamer
+import torch
+import torch.amp.autocast_mode
+from PIL import Image
+import torchvision.transforms.functional as TVF
+from threading import Thread
+from typing import Generator
+
+
+MODEL_PATH = "fancyfeast/llama-joycaption-alpha-two-vqa-test-1"
+TITLE = "<h1><center>JoyCaption Alpha Two - VQA Test - (2024-11-25a)</center></h1>"
+DESCRIPTION = """
+<div>
+<p>🚨🚨🚨 BY USING THIS SPACE YOU AGREE THAT YOUR QUERIES (but not images) <i>MAY</i> BE LOGGED AND COLLECTED ANONYMOUSLY 🚨🚨🚨</p>
+<p>🧪🧪🧪 This an experiment to see how well JoyCaption Alpha Two can learn to answer questions about images and follow instructions.
+I've only finetuned it on 600 examples, so it is highly experimental, very weak, broken, and volatile.  But for only training 600 examples,
+I thought it was performing surprisingly well and wanted to share. 🧪🧪🧪</p>
+<p>Unlike JoyCaption Alpha Two, you can ask this finetune questions about the image, like "What is he holding in his hand?", "Where might this be?",
+and "What are they doing?".  It can also follow instructions, like "Write me a poem about this image",
+"Write a caption but don't use any ambigious language, and make sure you mention that the image is from Instagram.", and
+"Output JSON with the following properties: 'skin_tone', 'hair_style', 'hair_length', 'clothing', 'background'." Remember that this was only finetuned on
+600 VQA/instruction examples, so it is _very_ limited right now.  Expect it to frequently fallback to its base behavior of just writing image descriptions.
+Expect accuracy to be lower.  Expect glitches.  Despite that, I've found that it will follow most queries I've tested it with, even outside its training,
+with enough coaxing and re-rolling.</p>
+<p>About the 🚨🚨🚨 above: this space will log all prompts sent to it.  The only thing this space logs is the text query; no images, no user data, etc.
+I cannot see what images you send, and frankly, I don't want to.  But knowing what kinds of instructions and queries users want JoyCaption to handle will
+help guide me in building JoyCaption's VQA dataset.  I've found out the hard way that almost all public VQA datasets are garbage and don't do a good job of
+training and exercising visual understanding.  Certainly not good enough to handle the complicated instructions that will allow JoyCaption users to guide and
+direct how JoyCaption writes descriptions and captions.  So I'm building my own dataset, that will be made public.  So, with peace and love, this space logs the text
+queries.  As always, the model itself is completely public and free to use outside of this space.  And, of course, I have no control nor access to what HuggingFace,
+which are graciously hosting this space, log.</p>
+</div>
+"""
+
+PLACEHOLDER = """
+"""
+
+
+
+# Load model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=True)
+assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Expected PreTrainedTokenizer, got {type(tokenizer)}"
+
+model = LlavaForConditionalGeneration.from_pretrained(MODEL_PATH, torch_dtype="bfloat16", device_map=0)
+assert isinstance(model, LlavaForConditionalGeneration), f"Expected LlavaForConditionalGeneration, got {type(model)}"
+
+
+def trim_off_prompt(input_ids: list[int], eoh_id: int, eot_id: int) -> list[int]:
+	# Trim off the prompt
+	while True:
+		try:
+			i = input_ids.index(eoh_id)
+		except ValueError:
+			break
+		
+		input_ids = input_ids[i + 1:]
+	
+	# Trim off the end
+	try:
+		i = input_ids.index(eot_id)
+	except ValueError:
+		return input_ids
+	
+	return input_ids[:i]
+
+end_of_header_id = tokenizer.convert_tokens_to_ids("<|end_header_id|>")
+end_of_turn_id = tokenizer.convert_tokens_to_ids("<|eot_id|>")
+assert isinstance(end_of_header_id, int) and isinstance(end_of_turn_id, int)
+
+
+@spaces.GPU()
+@torch.no_grad()
+def chat_joycaption(message: dict, history, temperature: float, max_new_tokens: int) -> Generator[str, None, None]:
+	torch.cuda.empty_cache()
+
+	# Prompts are always stripped in training for now
+	prompt = message['text'].strip()
+
+	# Load image
+	if "files" not in message or len(message["files"]) != 1:
+		raise ValueError("This model requires exactly one image as input.")
+	
+	image = Image.open(message["files"][0])
+	
+	# Log the prompt
+	print(f"Prompt: {prompt}")
+
+	# Preprocess image
+	# NOTE: I found the default processor for so400M to have worse results than just using PIL directly
+	if image.size != (384, 384):
+		image = image.resize((384, 384), Image.LANCZOS)
+	image = image.convert("RGB")
+	pixel_values = TVF.pil_to_tensor(image)
+
+	convo = [
+		{
+			"role": "system",
+			"content": "You are a helpful image captioner.",
+		},
+		{
+			"role": "user",
+			"content": prompt,
+		},
+	]
+
+	# Format the conversation
+	convo_string = tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = True)
+	assert isinstance(convo_string, str)
+
+	# Tokenize the conversation
+	convo_tokens = tokenizer.encode(convo_string, add_special_tokens=False, truncation=False)
+
+	# Repeat the image tokens
+	input_tokens = []
+	for token in convo_tokens:
+		if token == model.config.image_token_index:
+			input_tokens.extend([model.config.image_token_index] * model.config.image_seq_length)
+		else:
+			input_tokens.append(token)
+	
+	input_ids = torch.tensor(input_tokens, dtype=torch.long)
+	attention_mask = torch.ones_like(input_ids)
+
+	# Move to GPU
+	input_ids = input_ids.unsqueeze(0).to("cuda")
+	attention_mask = attention_mask.unsqueeze(0).to("cuda")
+	pixel_values = pixel_values.unsqueeze(0).to("cuda")
+
+	# Normalize the image
+	pixel_values = pixel_values / 255.0
+	pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
+	pixel_values = pixel_values.to(torch.bfloat16)
+
+	generate_kwargs = dict(
+		input_ids=input_ids,
+		pixel_values=pixel_values,
+		attention_mask=attention_mask,
+		max_new_tokens=max_new_tokens,
+		do_sample=True,
+		suppress_tokens=None,
+		use_cache=True,
+		temperature=temperature,
+		top_k=None,
+		top_p=0.9,
+		streamer=streamer,
+	)
+
+	if temperature == 0:
+		generate_kwargs["do_sample"] = False
+	
+	streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+	t = Thread(target=model.generate, kwargs=generate_kwargs)
+	t.start()
+
+	outputs = []
+	for text in streamer:
+		outputs.append(text)
+		yield "".join(outputs)
+
+
+chatbot=gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')
+
+with gr.Blocks() as demo:
+	gr.HTML(TITLE)
+	gr.Markdown(DESCRIPTION)
+	gr.ChatInterface(
+		fn=chat_joycaption,
+		chatbot=chatbot,
+		fill_height=True,
+		additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
+		additional_inputs=[
+			gr.Slider(minimum=0,
+						maximum=1, 
+						step=0.1,
+						value=0.6, 
+						label="Temperature", 
+						render=False),
+			gr.Slider(minimum=128, 
+						maximum=4096,
+						step=1,
+						value=1024, 
+						label="Max new tokens", 
+						render=False ),
+			],
+		examples=[
+			['How to setup a human base on Mars? Give short answer.'],
+			['Explain theory of relativity to me like I’m 8 years old.'],
+			['What is 9,000 * 9,000?'],
+			['Write a pun-filled happy birthday message to my friend Alex.'],
+			['Justify why a penguin might make a good king of the jungle.']
+			],
+		cache_examples=False,
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+huggingface_hub==0.23.4
+accelerate
+torch
+transformers==4.45.2
+sentencepiece
+torchvision