bczhou/clip-gpt2

1import gradio as gr2from clip_gpt2 import CLIPGPT2, CLIPGPT2Config, CLIPGPT2Processor3import os4import torch5 6device = 'cuda' if torch.cuda.is_available() else 'cpu'7 8config = CLIPGPT2Config(image_from_pretrained=False, text_from_pretrained=False)9model = CLIPGPT2(config)10model.load_state_dict(torch.load("pytorch_model.bin", map_location=device))11processor = CLIPGPT2Processor(config)12 13title = "Generate Image Captions With CLIP And GPT2"14 15 16def generate_image_captions(image, text):17    inputs = processor(images=image, texts=text, return_tensors="pt")18    input_ids = inputs.get("input_ids", None)19    pixel_values = inputs.get("pixel_values", None)20    attention_mask = inputs.get("attention_mask", None)21    prediction = model.generate(22        pixel_values=pixel_values,23        input_ids=input_ids,24        attention_mask=attention_mask,25        max_new_tokens=5026    )27    processor.tokenizer.padding_side = 'left'28    processor.tokenizer.pad_token_id = processor.tokenizer.eos_token_id29    prediction_text = processor.decode(prediction[0], skip_special_tokens=True)30    31    return prediction_text32 33article = "This demo is originated from this paper: [original paper](https://arxiv.org/abs/2209.15162)"34description = """35### Expand GPT2's language capabilities to vision with CLIP! 36### Tips:37- Only English is supported.38- When no image is provided, the model degrades to a vanilla GPT2-Large!39- When no description is provided, the model automatically generates a caption for the provided image.40- Try appending 'Answer:' after your question, the model is more likely to give desired outputs this way.41"""42demo = gr.Interface(43    fn=generate_image_captions,44    inputs=[45        gr.Image(),46        gr.Textbox(placeholder="A picture of", lines=3)47    ],48    outputs="text",49    examples=[50        [os.path.join(os.getcwd(), 'two_bear.png'), ""],51        [os.path.join(os.getcwd(), 'three_women.png'), "What is the woman in the middle's dress's color? Answer:"],52        [os.path.join(os.getcwd(), 'cat_with_food.png'), "Describe the picture:"],53        [os.path.join(os.getcwd(), 'dog_with_frisbee.png'), "What is the color of the frisbee in the photo? Answer:"],54        [os.path.join(os.getcwd(), 'stop_sign.png'), "What does the sign in the picture say? Answer:"]55    ],56    article=article,57    title=title,58    description=description,59    cache_examples=False60)61 62demo.launch()