Update app.py

app.py

@@ -3,112 +3,74 @@ from PIL import Image
 import streamlit as st
 import os
 import random
-import numpy as np
 import torch
-from diffusers import FluxPipeline  # Importar FluxPipeline
+from diffusers import StableDiffusionImg2ImgPipeline  # Reemplazar FluxPipeline
 
-# Configuraciones y cliente de inferencia
-MAX_SEED = np.iinfo(np.int32).max
+# Configuraciones y ruta de guardado
 DATA_PATH = Path("./data")
 DATA_PATH.mkdir(exist_ok=True)
 
-# Cargar modelo de Flux
+# Cargar modelo de Stable Diffusion
 device = "cuda" if torch.cuda.is_available() else "cpu"
-flux_pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16).to(device)
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
 
-# Generación de imagen usando prompt y/o imagen subida
-async def gen(prompt, width, height, uploaded_image=None):
-    images = []
-    try:
-        seed = random.randint(0, MAX_SEED)
-        if uploaded_image:
-            uploaded_image = Image.open(uploaded_image).convert("RGB")
-            image = await generate_img2img(prompt, uploaded_image)
-        else:
-            image = await generate_image(prompt, width, height, seed)
-        image_path = save_image(image, f"generated_image_{seed}.jpg", prompt)
-        if image_path:
-            images.append(str(image_path))
-    except Exception as e:
-        st.error(f"Error al generar imágenes: {e}")
-    return images
+# Generar imagen con Stable Diffusion (txt2img o img2img)
+def generate_image(prompt, init_image=None, strength=0.75):
+    generator = torch.manual_seed(random.randint(0, 1000000))
+    if init_image:
+        image = pipe(prompt=prompt, init_image=init_image, strength=strength, generator=generator).images[0]
+    else:
+        image = pipe(prompt=prompt, generator=generator).images[0]
+    return image
+
+# Guardar imagen
+def save_image(image, file_name, prompt=None):
+    image_path = DATA_PATH / file_name
+    image.save(image_path, format="JPEG")
+    if prompt:
+        with open(DATA_PATH / "prompts.txt", "a") as f:
+            f.write(f"{file_name}: {prompt}\n")
+    return image_path
 
-# Guardar y mostrar galería
+# Mostrar galería
 def list_saved_images():
     return list(DATA_PATH.glob("*.jpg"))
 
 def display_gallery():
-    st.header("Galería de Imágenes Guardadas")
+    st.header("Galería de Imágenes")
     images = list_saved_images()
     if images:
-        cols = st.columns(8)
+        cols = st.columns(4)
         for i, image_file in enumerate(images):
-            with cols[i % 8]:
+            with cols[i % 4]:
                 st.image(str(image_file), caption=image_file.name, use_column_width=True)
-                prompt = get_prompt_for_image(image_file.name)
-                st.write(prompt[:300])
-
-                if st.button(f"Borrar", key=f"delete_{i}_{image_file.name}"):
+                if st.button(f"Borrar {image_file.name}", key=f"delete_{i}"):
                     os.remove(image_file)
-                    st.success("Imagen borrada")
-                    display_gallery()
+                    st.success(f"Imagen {image_file.name} borrada.")
     else:
         st.info("No hay imágenes guardadas.")
 
-def save_image(image, file_name, prompt=None):
-    image_path = DATA_PATH / file_name
-    if image_path.exists():
-        st.warning(f"La imagen '{file_name}' ya existe en la galería. No se guardó.")
-        return None
-    else:
-        image.save(image_path, format="JPEG")
-        if prompt:
-            save_prompt(f"{file_name}: {prompt}")
-        return image_path
-
-# Generación de imagen desde texto (txt2img)
-async def generate_image(prompt, width, height, seed):
-    image = flux_pipeline(prompt=prompt, width=width, height=height, num_inference_steps=4).images[0]
-    return image
-
-# Generación de imagen desde imagen y texto (img+txt=img)
-async def generate_img2img(prompt, image):
-    image = flux_pipeline(prompt=prompt, init_image=image, num_inference_steps=4).images[0]
-    return image
-
-def get_prompt_for_image(image_name):
-    prompts = {}
-    try:
-        with open(DATA_PATH / "prompts.txt", "r") as f:
-            for line in f:
-                if line.startswith(image_name):
-                    prompts[image_name] = line.split(": ", 1)[1].strip()
-    except FileNotFoundError:
-        return "No hay prompt asociado."
-    return prompts.get(image_name, "No hay prompt asociado.")
-
-# Función principal
-async def main():
+# Función principal de la app
+def main():
     st.set_page_config(layout="wide")
+    st.title("Generación de Imágenes - Stable Diffusion")
 
-    st.title("Flux +Image +Variants")
-    prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=500)
-    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
-    width, height = (360, 640) if format_option == "9:16" else (640, 360)
-
-    uploaded_image = st.sidebar.file_uploader("Sube una imagen (opcional)", type=["png", "jpg", "jpeg"])
+    prompt = st.sidebar.text_area("Descripción de la imagen", height=150)
+    uploaded_image = st.sidebar.file_uploader("Sube una imagen para img2img (opcional)", type=["png", "jpg", "jpeg"])
 
     if st.sidebar.button("Generar Imagen"):
-        with st.spinner("Generando imágenes..."):
-            try:
-                results = await gen(prompt, width, height, uploaded_image)
-                for result in results:
-                    st.image(result, caption="Imagen Generada")
-            except Exception as e:
-                st.error(f"Error al generar las imágenes: {str(e)}")
+        with st.spinner("Generando imagen..."):
+            if uploaded_image:
+                init_image = Image.open(uploaded_image).convert("RGB")
+                generated_image = generate_image(prompt, init_image)
+            else:
+                generated_image = generate_image(prompt)
+
+            image_path = save_image(generated_image, f"generated_{random.randint(0, 10000)}.jpg", prompt)
+            st.image(str(image_path), caption="Imagen Generada")
 
     display_gallery()
 
-# Ejecución del código principal
+# Ejecución de la app
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()