File size: 2,040 Bytes
e023f7e db58a82 e023f7e d8b0f0c e023f7e db58a82 e023f7e db58a82 |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import inspect
import gradio as gr
import torch
from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
from PIL import Image
class AestheticPredictor:
def __init__(self):
# load model and preprocessor
self.model, self.preprocessor = convert_v2_5_from_siglip(
low_cpu_mem_usage=True,
trust_remote_code=True,
)
if torch.cuda.is_available():
self.model = self.model.to(torch.bfloat16).cuda()
def inference(self, image: Image.Image) -> float:
# preprocess image
pixel_values = self.preprocessor(
images=image.convert("RGB"), return_tensors="pt"
).pixel_values
if torch.cuda.is_available():
pixel_values = pixel_values.to(torch.bfloat16).cuda()
# predict aesthetic score
with torch.inference_mode():
score = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
return score
if __name__ == "__main__":
aesthetic_predictor = AestheticPredictor()
with gr.Blocks(theme="soft") as blocks:
markdown = gr.Markdown(
value=inspect.cleandoc(
"""
# Aesthetic Predictor V2.5
This app predicts the aesthetic score of input images such as paintings,
photographs, and illustrations.
The aesthetic score is a floating-point number between 1 and 10.
5.5+ is considered to be a good aesthetic score.
You can get the model from [Github](https://github.com/discus0434/aesthetic-predictor-v2-5?tab=readme-ov-file).
"""
)
)
with gr.Row():
with gr.Column():
image = gr.Image(label="Input Image", type="pil")
button = gr.Button("Predict")
with gr.Column():
score = gr.Textbox(label="Aesthetic Score")
button.click(aesthetic_predictor.inference, inputs=image, outputs=score)
blocks.queue().launch()
|