import streamlit as st
from huggingface_hub import InferenceClient
from PIL import Image
import io

# Initialize Hugging Face Inference Client
hf_token = st.secrets["HF_TOKEN"]  # Use your Hugging Face token
client = InferenceClient(model="enhanceaiteam/Flux-uncensored", token=hf_token)

# Streamlit UI
st.title("Generación de Imágenes img2img con Flux")
st.sidebar.header("Opciones de generación")
prompt = st.sidebar.text_input("Escribe tu prompt:", value="")

uploaded_file = st.sidebar.file_uploader("Sube una imagen base", type=["jpg", "png"])

if uploaded_file is not None:
    try:
        init_image = Image.open(uploaded_file).convert("RGB")
        st.image(init_image, caption="Imagen Base", use_column_width=True)

        if st.sidebar.button("Generar Imagen"):
            with st.spinner('Generando imagen...'):
                # Convert the uploaded image to bytes
                img_bytes = io.BytesIO()
                init_image.save(img_bytes, format="PNG")
                img_bytes = img_bytes.getvalue()

                # Call the API with the prompt and base image
                image = client.img2img(prompt=prompt, image=img_bytes)  # Returns image bytes

                # Display the generated image
                st.image(image, caption="Imagen Generada")

                # Option to download the generated image
                st.download_button(label="Descargar imagen", data=image, file_name="generated_image.png", mime="image/png")

    except Exception as e:
        st.error(f"Error en la generación de la imagen: {e}")
else:
    st.warning("Por favor, sube una imagen base para continuar.")