Update app.py

app.py

@@ -2,14 +2,27 @@ import streamlit as st
 import torch
 from PIL import Image
 import random
-from diffusers import StableDiffusionInstructPix2PixPipeline
+from diffusers import StableDiffusionInstructPix2PixPipeline, DiffusionPipeline, FluxPipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-selected_model = "timbrooks/instruct-pix2pix"
+
+model_options = {
+    "Instruct Pix2Pix": "timbrooks/instruct-pix2pix",
+    "SDXL Turbo": "stabilityai/sdxl-turbo",
+    "FLUX": "alimama-creative/FLUX.1-Turbo-Alpha"
+}
+
+selected_model = st.sidebar.selectbox("Selecciona el modelo", list(model_options.keys()))
+model_id = model_options[selected_model]
 
 @st.cache_resource
 def load_model():
-    return StableDiffusionInstructPix2PixPipeline.from_pretrained(selected_model, torch_dtype=torch.float32, safety_checker=None).to(device)
+    if selected_model == "Instruct Pix2Pix":
+        return StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float32, safety_checker=None).to(device)
+    elif selected_model == "SDXL Turbo":
+        return DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, safety_checker=None).to(device)
+    else:
+        return FluxPipeline.from_pretrained(model_id, torch_dtype=torch.float32, safety_checker=None).to(device)
 
 pipe = load_model()
 
@@ -26,17 +39,34 @@ def infer(source_img, prompt, steps, seed, text_cfg_scale, image_cfg_scale):
 
     try:
         for i in range(steps):
-            output = pipe(
-                prompt,
-                image=source_image,
-                guidance_scale=text_cfg_scale,
-                num_inference_steps=1,
-                image_guidance_scale=image_cfg_scale,
-                safety_checker=None  # Desactivar el verificador de seguridad
-            )
-            if len(output.images) == 0:
+            if selected_model == "Instruct Pix2Pix":
+                output = pipe(
+                    prompt,
+                    image=source_image,
+                    guidance_scale=text_cfg_scale,
+                    num_inference_steps=1,
+                    image_guidance_scale=image_cfg_scale
+                )
+            elif selected_model == "SDXL Turbo":
+                output = pipe(
+                    prompt,
+                    num_inference_steps=steps,
+                    guidance_scale=text_cfg_scale,
+                    image=source_image
+                )
+            else:
+                output = pipe(
+                    prompt,
+                    image=source_image,
+                    num_inference_steps=steps,
+                    guidance_scale=text_cfg_scale
+                )
+
+            result = output.images[0] if output.images else None
+            
+            if result is None:
                 raise ValueError("No se generaron imágenes.")
-            result = output.images[0]
+
             progress_bar.progress((i + 1) / steps) 
 
         progress_bar.progress(1.0)
@@ -48,19 +78,17 @@ def infer(source_img, prompt, steps, seed, text_cfg_scale, image_cfg_scale):
 
 st.title("Flux Image to Image")
 
-# Sidebar para controles
 with st.sidebar:
     uploaded_image = st.file_uploader("Sube una imagen", type=["png", "jpg", "jpeg"], key="unique_file_uploader")
     prompt = st.text_input("Texto del prompt (máx. 77 tokens)")
-    steps = st.slider("Número de Iteraciones", min_value=1, max_value=50, value=2, step=1)
-    randomize_seed = st.radio("Randomize Seed", ["Fix Seed", "Randomize Seed"])
+    steps = st.slider("Número de Iteraciones", min_value=1, max_value=50, value=11, step=1)
+    randomize_seed = st.radio("Randomize Seed", ["Randomize Seed", "Fix Seed"])
     seed = st.slider("Seed", min_value=0, max_value=9999, step=1, value=random.randint(0, 9999) if randomize_seed == "Randomize Seed" else 1111)
-    text_cfg_scale = st.slider("Text CFG Scale", min_value=0.0, max_value=10.0, step=0.1, value=7.5)
-    image_cfg_scale = st.slider("Image CFG Scale", min_value=0.0, max_value=10.0, step=0.1, value=1.5)
+    text_cfg_scale = st.slider("Text CFG Scale", min_value=0.0, max_value=10.0, step=0.1, value=9.0)
+    image_cfg_scale = st.slider("Image CFG Scale", min_value=0.0, max_value=10.0, step=0.1, value=1.0)
 
-# Columna para mostrar la imagen generada
 if uploaded_image is not None and st.button("Generar imagen"):
-    image = Image.open(uploaded_image)
+    image = Image.open(uploaded_image).convert("RGB")
     result_image = infer(image, prompt, steps, seed, text_cfg_scale, image_cfg_scale)
     if result_image is not None:
-        st.image(result_image, caption="Imagen generada", use_column_width=True)  # Mostrar la imagen generada
+        st.image(result_image, caption="Imagen generada", use_column_width=True)