Update app.py

app.py

@@ -1,41 +1,47 @@
-import streamlit as st
+import base64
+import io
 from huggingface_hub import InferenceClient
+import streamlit as st
 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)
+# Configuración
+hf_token = st.secrets["HF_TOKEN"]
+client = InferenceClient(repo_id="black-forest-labs/FLUX.1-schnell", token=hf_token)
+
+def get_image_result(prompt, init_image=None):
+    if init_image:
+        img_buffer = io.BytesIO()
+        init_image.save(img_buffer, format="PNG")
+        img_bytes = img_buffer.getvalue()
+        encoded_image = base64.b64encode(img_bytes).decode('utf-8')
+    else:
+        encoded_image = None
+
+    try:
+        output = client.img2img(prompt=prompt, image=encoded_image, strength=0.75)
+        img_data = base64.b64decode(output['generated_image'])
+        return Image.open(io.BytesIO(img_data))
+    except Exception as e:
+        st.error(f"Error en la generación de la imagen: {str(e)}")
+        return None
 
-# Streamlit UI
+# Interfaz de Streamlit
 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="")
+prompt = st.sidebar.text_input("Escribe tu prompt:", value="Una escena futurista en una ciudad verde")
+uploaded_file = st.sidebar.file_uploader("Sube una imagen base (opcional)", type=["jpg", "png"])
 
-uploaded_file = st.sidebar.file_uploader("Sube una imagen base", type=["jpg", "png"])
-
-if uploaded_file is not None:
-    try:
+if st.sidebar.button("Generar Imagen"):
+    if uploaded_file:
         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.")
+    else:
+        init_image = None
+
+    with st.spinner('Generando imagen...'):
+        generated_image = get_image_result(prompt, init_image)
+        if generated_image:
+            st.image(generated_image, caption="Imagen Generada", use_column_width=True)
+            img_bytes = io.BytesIO()
+            generated_image.save(img_bytes, format="PNG")
+            st.download_button(label="Descargar imagen", data=img_bytes.getvalue(), file_name="generated_image.png", mime="image/png")