Update app.py

app.py

@@ -18,11 +18,11 @@ model_id = model_options[selected_model]
 @st.cache_resource
 def load_model():
     if selected_model == "Instruct Pix2Pix":
-        return StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float32, safety_checker=None).to(device)
+        return StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, 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)
+        return DiffusionPipeline.from_pretrained(model_id, safety_checker=None).to(device)
     else:
-        return FluxPipeline.from_pretrained(model_id, torch_dtype=torch.float32, safety_checker=None).to(device)
+        return FluxPipeline.from_pretrained(model_id, safety_checker=None).to(device)
 
 pipe = load_model()
 
@@ -31,7 +31,7 @@ def resize(img, max_size=512):
     new_size = (int(img.width * ratio), int(img.height * ratio))
     return img.resize(new_size, Image.LANCZOS)
 
-def infer(source_img, prompt, steps, seed, text_cfg_scale, image_cfg_scale):
+def infer(source_img, prompt, steps, seed, guidance_scale):
     torch.manual_seed(seed)
     source_image = resize(source_img)
     progress_bar = st.progress(0)
@@ -39,35 +39,19 @@ def infer(source_img, prompt, steps, seed, text_cfg_scale, image_cfg_scale):
 
     try:
         for i in range(steps):
-            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
-                )
+            output = pipe(
+                prompt,
+                image=source_image,
+                guidance_scale=guidance_scale / 10,  
+                num_inference_steps=steps
+            )
 
             result = output.images[0] if output.images else None
-            
+
             if result is None:
                 raise ValueError("No se generaron imágenes.")
 
-            progress_bar.progress((i + 1) / steps) 
+            progress_bar.progress((i + 1) / steps)
 
         progress_bar.progress(1.0)
     except Exception as e:
@@ -84,11 +68,10 @@ with st.sidebar:
     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=9.0)
-    image_cfg_scale = st.slider("Image CFG Scale", min_value=0.0, max_value=10.0, step=0.1, value=1.0)
+    guidance_scale = st.slider("Guidance Scale", min_value=0.0, max_value=10.0, step=0.01, value=9.0)
 
 if uploaded_image is not None and st.button("Generar imagen"):
     image = Image.open(uploaded_image).convert("RGB")
-    result_image = infer(image, prompt, steps, seed, text_cfg_scale, image_cfg_scale)
+    result_image = infer(image, prompt, steps, seed, guidance_scale)
     if result_image is not None:
-        st.image(result_image, caption="Imagen generada", use_column_width=True)
+        st.image(result_image, caption="Imagen generada", use_column_width=True)