File size: 1,004 Bytes
d1b0574
c65a1f7
 
 
5a4a035
d1b0574
ea1edb1
 
 
382da0b
c65a1f7
0e2f2fd
 
 
c65a1f7
cfc2299
1506268
 
 
 
 
 
 
c65a1f7
 
d1b0574
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import gradio as gr
from transformers import AutoModel, AutoProcessor
from PIL import Image
import torch
import numpy as np

model_name_or_path = "lyttt/VLV_captioner"
model = AutoModel.from_pretrained(model_name_or_path, revision="master", trust_remote_code=True,low_cpu_mem_usage=False)

# @spaces.GPU(duration=120)
def greet(image):
    if image.dtype != np.uint8:
        image = (np.clip(image, 0, 1) * 255).astype(np.uint8)
    image = Image.fromarray(image, mode='RGB')
    with torch.no_grad():
        outputs = model([image], 300).generated_text[0]
    def drop_incomplete_tail(text):
        sentences = text.split('.')
        complete_sentences = [s.strip() for s in sentences if s.strip()]
        if not text.strip().endswith('.'):
            complete_sentences = complete_sentences[:-1]
        return '. '.join(complete_sentences) + ('.' if complete_sentences else '')
    return drop_incomplete_tail(outputs)

demo = gr.Interface(fn=greet, inputs="image", outputs="text")
demo.launch()