1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.image.image_aesthetic_score import ImageAestheticScore
10
11if __name__ == "__main__":
12    if os.getenv("DAFT_RUNNER", "native") == "ray":
13        import ray
14        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15        daft.set_runner_ray()
16    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18    # TODO: 根据实际场景准备样本数据
19    samples = {"image_inputs": [...]}
20    ds = daft.from_pydict(samples)
21    constructor_kwargs = {
22        "batch_size": 32,
23        "model_path": '/opt/aihc/models',
24        "clip_model_name": 'openai/clip-vit-large-patch14',
25        "mlp_model_name": 'laion_aesthetic_v2/sac+logos+ava1-l14-linearMSE.pth',
26        "device": 'cpu',
27    }
28    ds = ds.with_column(
29        "result",
30        aihc_udf(
31            ImageAestheticScore,
32            construct_args=constructor_kwargs,
33            num_cpus=1,
34            concurrency=4,
35            batch_size=8,
36        )(col("image_inputs")),
37    )
38    ds.show()