app.py

# type: ignore
from typing import Any
import gradio as gr
import spaces
import torch
from PIL import Image
from transformers import AutoModelForCausalLM, LlamaTokenizer


DEFAULT_PARAMS = {
 "do_sample": False,
 "max_new_tokens": 256,
}
DEFAULT_QUERY = (
 "Provide a factual description of this image in up to two paragraphs. "
 "Include details on objects, background, scenery, interactions, gestures, poses, and any visible text content. "
 "Specify the number of repeated objects. "
 "Describe the dominant colors, color contrasts, textures, and materials. "
 "Mention the composition, including the arrangement of elements and focus points. "
 "Note the camera angle or perspective, and provide any identifiable contextual information. "
 "Include details on the style, lighting, and shadows. "
 "Avoid subjective interpretations or speculation."
)

DTYPE = torch.bfloat16
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = LlamaTokenizer.from_pretrained(
 pretrained_model_name_or_path="lmsys/vicuna-7b-v1.5",
)
model = AutoModelForCausalLM.from_pretrained(
 pretrained_model_name_or_path="THUDM/cogvlm-chat-hf",
 torch_dtype=DTYPE,
 trust_remote_code=True,
 low_cpu_mem_usage=True,
)

model = model.to(device=DEVICE)

@spaces.GPU
@torch.no_grad()
def generate_caption(
 image: Image.Image,
 query: str = DEFAULT_QUERY,
 params: dict[str, Any] = DEFAULT_PARAMS,
) -> str:
 inputs = model.build_conversation_input_ids(
 tokenizer=tokenizer,
 query=query,
 history=[],
 images=[image],
 )
 inputs = {
 "input_ids": inputs["input_ids"].unsqueeze(0).to(device=DEVICE),
 "token_type_ids": inputs["token_type_ids"].unsqueeze(0).to(device=DEVICE),
 "attention_mask": inputs["attention_mask"].unsqueeze(0).to(device=DEVICE),
 "images": [[inputs["images"][0].to(device=DEVICE, dtype=DTYPE)]],
 }

 outputs = model.generate(**inputs, **params)
 outputs = outputs[:, inputs["input_ids"].shape[1] :]
 result = tokenizer.decode(outputs[0])

 result = result.replace("This image showcases", "").strip().removesuffix("</s>").strip().capitalize()
 return result


with gr.Blocks() as demo:
 with gr.Row():
 with gr.Column():
 input_image = gr.Image(type="pil")
 input_query = gr.Textbox(lines=5, label="Prompt", value=DEFAULT_QUERY)
 run_button = gr.Button(value="Generate Caption")
 with gr.Column():
 output_caption = gr.Textbox(label="Generated Caption", show_copy_button=True)

 run_button.click(
 fn=generate_caption,
 inputs=[input_image, input_query],
 outputs=output_caption,
 )

demo.launch(share=False)