import streamlit as st
from PIL import Image
import requests
from io import BytesIO
import os

# 환경 변수에서 API 토큰을 가져옵니다.
api_token = os.getenv("HUGGINGFACE_API_TOKEN")
if not api_token:
    st.error("API token not found. Please set the HUGGINGFACE_API_TOKEN environment variable.")

def transform_image(image):
    api_url = "https://api-inference.huggingface.co/models/akhaliq/AnimeGANv2"
    headers = {"Authorization": f"Bearer {api_token}"}

    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    buffered.seek(0)

    files = {"file": ("image.jpg", buffered, "image/jpeg")}
    response = requests.post(api_url, headers=headers, files=files)

    if response.status_code == 200:
        return Image.open(BytesIO(response.content))
    else:
        st.error(f"Failed to transform image. Status code: {response.status_code}. Response: {response.text}")
        return None

st.title("AnimeGANv2 Image Transformer")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)

    if st.button("Transform"):
        transformed_image = transform_image(image)
        if transformed_image:
            st.image(transformed_image, caption='Transformed Image.', use_column_width=True)