worker.py

from celery_app import AI_detection_celery_app
from web_inference import do_web_inference
from model_with_pure_bert import MusicAudioClassifier
from ISMIR_2025.MERT.networks import MERTFeatureExtractor


checkpoint_path = "with_embedding_MERT_768_embedding/EmbeddingModel_MERT_768-epoch=0353-val_loss=0.3866-val_acc=0.9809-val_f1=0.9803-val_precision=0.9764-val_recall=0.9842.ckpt"  # 실제 경로로 조정
backbone_model = MERTFeatureExtractor()
input_dim = 768
# 모델 로드
model = MusicAudioClassifier.load_from_checkpoint(
    checkpoint_path, 
    input_dim=input_dim, 
    emb_model=backbone_model,
    strict=False   
)

@AI_detection_celery_app.task(name='AI_detection_task')
def AI_detection_task(data):
    music_id =  data.get('music_id')
    do_web_inference(model, music_id)
    return {"status":"done"}