File size: 1,258 Bytes
d93f45c
b51ba4d
7852f50
 
 
 
 
 
 
 
b51ba4d
d93f45c
b51ba4d
 
 
 
bea1510
d93f45c
7852f50
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
27
28
29
30
31
32
33
34
35
36
37
from fastapi import FastAPI
from fastapi.responses import RedirectResponse
# from transformers import pipeline
from PIL import Image
import requests
from transformers import AutoProcessor, Pix2StructForConditionalGeneration

processor = AutoProcessor.from_pretrained("google/pix2struct-textcaps-base")
model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-textcaps-base")


app = FastAPI()

@app.get("/")
async def docs_redirect():
    return RedirectResponse(url='/docs')

@app.get("/generate")
def generate():
    url = "https://www.ilankelman.org/stopsigns/australia.jpg"
    image = Image.open(requests.get(url, stream=True).raw)

    inputs = processor(images=image, return_tensors="pt")

    # autoregressive generation
    generated_ids = model.generate(**inputs, max_new_tokens=50)
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(generated_text)

    # conditional generation
    text = "A picture of"
    inputs = processor(text=text, images=image, return_tensors="pt", add_special_tokens=False)

    generated_ids = model.generate(**inputs, max_new_tokens=50)
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(generated_text)