|
import logging |
|
import os |
|
|
|
import gradio as gr |
|
from pillow_heif import register_heif_opener |
|
|
|
register_heif_opener() |
|
|
|
import gradio as gr |
|
from transformers import pipeline |
|
|
|
LOG_LEVEL = os.getenv("LOG_LEVEL", "DEBUG") |
|
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 200)) |
|
|
|
MODEL = os.getenv("MODEL", "Salesforce/blip-image-captioning-large") |
|
|
|
logging.basicConfig(level=LOG_LEVEL) |
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
logger.info("Loading model...") |
|
|
|
captioner = pipeline( |
|
"image-to-text", |
|
model=MODEL, |
|
max_new_tokens=MAX_NEW_TOKENS, |
|
) |
|
logger.info("Done loading model.") |
|
|
|
|
|
def graptioner(image_url): |
|
global captioner |
|
result = captioner(image_url) |
|
caption = result[0]["generated_text"] |
|
return caption |
|
|
|
|
|
|
|
iface = gr.Interface(fn=graptioner, inputs="text", outputs=["text"], allow_flagging="never") |
|
iface.launch() |
|
|