# 0. 필수 라이브러리 설치 (Colab에서 실행 시 주석 해제 후 먼저 실행)
# !pip install gradio transformers torch pillow

import gradio as gr
import torch
import spaces
from PIL import Image
from transformers import CLIPProcessor, CLIPModel

# 1. 모델 및 프로세서 로드
# 모델 다운로드를 위해 최초 실행 시 시간이 조금 걸릴 수 있습니다.
model_name = "openai/clip-vit-base-patch32"
try:
    model = CLIPModel.from_pretrained(model_name)
    processor = CLIPProcessor.from_pretrained(model_name)
    print(f"모델 로드 완료: {model_name}")
except Exception as e:
    print(f"모델 로드 중 오류 발생: {e}")

# GPU 사용 가능 시 CUDA 설정, 아니면 CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# 2. 추론 함수 정의
@spaces.GPU
def predict(image, text_options):
    if image is None:
        return None
    
    # 텍스트 전처리 (쉼표로 구분된 문자열을 리스트로 변환하고 공백 제거)
    candidates = [t.strip() for t in text_options.split(",") if t.strip()]
    
    if not candidates:
        return {"Error": "텍스트 후보를 입력해주세요."}

    # 입력 데이터 처리 (이미지 + 텍스트)
    inputs = processor(
        text=candidates, 
        images=image, 
        return_tensors="pt", 
        padding=True
    ).to(device)

    # 추론 (Gradient 계산 불필요)
    with torch.no_grad():
        outputs = model(**inputs)

    # 결과 계산 (Logits -> Softmax 확률 변환)
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1).cpu().numpy()[0]

    # 결과 딕셔너리 생성 (Label: Score 형태)
    # Gradio의 Label 컴포넌트는 {라벨: 확률} 형태의 딕셔너리를 받습니다.
    return {candidates[i]: float(probs[i]) for i in range(len(candidates))}

# 3. Gradio 인터페이스 구성
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="이미지 업로드"),
        gr.Textbox(
            label="후보 텍스트 (쉼표로 구분)", 
            placeholder="예: soccer player, baseball player, referee", # 예시 수정
            value="cat, dog, car" # 기본값 설정
        )
    ],
    outputs=gr.Label(num_top_classes=3, label="매칭 결과"),
    title="CLIP 이미지-텍스트 매칭기",
    description="이미지를 업로드하고, 그 이미지가 무엇인지 설명하는 단어들을 쉼표(,)로 구분해 적어주세요. AI가 가장 적절한 설명을 찾아줍니다."
)

# 메인 실행 블록 수정 (오타 수정: _= -> ==)
if __name__ == "__main__":
    # Colab에서 실행 시 share=True를 하면 외부에서 접속 가능한 링크가 생성됩니다.
    iface.launch(share=True)