import gradio as gr
from transformers import pipeline
import requests
import os

# 환경변수에서 Hugging Face API 토큰을 로드
hugging_face_auth_token = os.getenv("HUGGING_FACE_AUTH_TOKEN")

# 이미지 인식 파이프라인 로드
image_model = pipeline("image-classification", model="google/vit-base-patch16-224")

def get_audiogen(prompt):
    # 오디오 생성 모델 API 호출
    response = requests.post(
        "https://api-inference.huggingface.co/models/fffiloni/audiogen",
        headers={"Authorization": f"Bearer {hugging_face_auth_token}"},  # 수정된 부분
        json={"inputs": prompt, "parameters": {"length": 10}, "options": {"use_cache": False}}
    )
    result = response.json()
    # 생성된 오디오 파일의 URL을 반환하거나, 오디오 데이터 자체를 반환
    return result

def classify_and_generate_audio(uploaded_image):
    # 이미지 분류
    predictions = image_model(uploaded_image)
    top_prediction = predictions[0]['label']  # 가장 확률이 높은 분류 결과
    
    # 오디오 생성
    audio_result = get_audiogen(top_prediction)
    
    # audio_result를 처리하여 Gradio가 재생할 수 있는 형식으로 반환
    return top_prediction, audio_result

# Gradio 인터페이스 생성
iface = gr.Interface(
    fn=classify_and_generate_audio,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Label(), gr.Audio()],
    title="이미지 분류 및 오디오 생성",
    description="이미지를 업로드하면, 이미지를 분석하여 무엇인지 설명하고, 해당하는 오디오를 생성합니다."
)

# 인터페이스 실행
iface.launch()