Update app.py

app.py

@@ -9,13 +9,16 @@ from gradio_client import Client, handle_file
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 import yaml
+import cv2
+import dlib
 
+# Cargar configuración
 try:
     with open("config.yaml", "r") as file:
         credentials = yaml.safe_load(file)
 except Exception as e:
     st.error(f"Error al cargar el archivo de configuración: {e}")
-    credentials = {"username": "", "password": ""} 
+    credentials = {"username": "", "password": ""}
 
 MAX_SEED = np.iinfo(np.int32).max
 HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
@@ -32,10 +35,8 @@ def run_async(func):
     return loop.run_until_complete(result)
 
 async def generate_image(combined_prompt, model, width, height, scales, steps, seed):
+    seed = int(seed) if seed != -1 else random.randint(0, MAX_SEED)
     try:
-        if seed == -1:
-            seed = random.randint(0, MAX_SEED)
-        seed = int(seed)
         image = await client.text_to_image(
             prompt=combined_prompt, height=height, width=width, guidance_scale=scales,
             num_inference_steps=steps, model=model
@@ -51,7 +52,7 @@ def get_upscale_finegrain(prompt, img_path, upscale_factor):
             input_image=handle_file(img_path), prompt=prompt, upscale_factor=upscale_factor
         )
         return result[1] if isinstance(result, list) and len(result) > 1 else None
-    except Exception as e:
+    except Exception:
         return None
 
 def save_prompt(prompt_text, seed):
@@ -65,15 +66,10 @@ def save_prompt(prompt_text, seed):
         return None
 
 async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language):
-    combined_prompt = prompt  
-    if process_enhancer:
-        improved_prompt = await improve_prompt(prompt, language)
-        combined_prompt = f"{prompt} {improved_prompt}"
-    
-    if seed == -1:
-        seed = random.randint(0, MAX_SEED)
-    seed = int(seed)
+    combined_prompt = f"{prompt} {await improve_prompt(prompt, language) if process_enhancer else ''}".strip()
+    seed = int(seed) if seed != -1 else random.randint(0, MAX_SEED)
     progress_bar = st.progress(0)
+
     image, seed = await generate_image(combined_prompt, basemodel, width, height, scales, steps, seed)
     progress_bar.progress(50)
 
@@ -87,29 +83,26 @@ async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_fac
     if process_upscale:
         upscale_image_path = get_upscale_finegrain(combined_prompt, image_path, upscale_factor)
         if upscale_image_path:
-            upscale_image = Image.open(upscale_image_path)
-            upscale_image.save(DATA_PATH / f"upscale_image_{seed}.jpg", format="JPEG")
+            Image.open(upscale_image_path).save(DATA_PATH / f"upscale_image_{seed}.jpg", format="JPEG")
             progress_bar.progress(100)
-            image_path.unlink()  
+            image_path.unlink()
             return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
-        else:
-            progress_bar.empty()
-            return [str(image_path), str(prompt_file_path)]
-    else:
-        progress_bar.progress(100)
-        return [str(image_path), str(prompt_file_path)]
+    
+    progress_bar.progress(100)
+    return [str(image_path), str(prompt_file_path)]
 
 async def improve_prompt(prompt, language):
-    try:
-        instruction_en = "With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add illumination, atmosphere, cinematic elements, and characters if needed..."
-        instruction_es = "Con esta idea, describe en español un prompt detallado de txt2img en un máximo de 500 caracteres, con iluminación, atmósfera, elementos cinematográficos y en su caso personajes..."
-        instruction = instruction_en if language == "en" else instruction_es
-        formatted_prompt = f"{prompt}: {instruction}"
-        response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
-        improved_text = response['generated_text'].strip() if 'generated_text' in response else response.strip()
-        return improved_text[:500] if len(improved_text) > 500 else improved_text
-    except Exception as e:
-        return f"Error mejorando el prompt: {e}"
+    instruction = (
+        "Con esta idea, describe en español un prompt detallado de txt2img en un máximo de 500 caracteres, "
+        "con iluminación, atmósfera, elementos cinematográficos y en su caso personajes..."
+        if language == "es" else
+        "With this idea, describe in English a detailed txt2img prompt in 500 characters at most, "
+        "add illumination, atmosphere, cinematic elements, and characters if needed..."
+    )
+    formatted_prompt = f"{prompt}: {instruction}"
+    response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
+    improved_text = response.get('generated_text', '').strip() if 'generated_text' in response else response.strip()
+    return improved_text[:500] if len(improved_text) > 500 else improved_text
 
 def save_image(image, seed):
     try:
@@ -123,8 +116,8 @@ def save_image(image, seed):
 def get_storage():
     files = [file for file in DATA_PATH.glob("*.jpg") if file.is_file()]
     files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
-    usage = sum([file.stat().st_size for file in files])
-    return [str(file.resolve()) for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"
+    usage = sum(file.stat().st_size for file in files)
+    return [str(file.resolve()) for file in files], f"Uso total: {usage / (1024.0 ** 3):.3f}GB"
 
 def get_prompts():
     prompt_files = [file for file in DATA_PATH.glob("*.txt") if file.is_file()]
@@ -142,8 +135,31 @@ def delete_image(image_path):
 
 def swap_faces(image_path):
     try:
-        swapped_image_path = image_path  # Simulación del swap; aquí debes implementar la lógica real.
-        return swapped_image_path
+        image = cv2.imread(str(image_path))
+        detector = dlib.get_frontal_face_detector()
+        predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+        faces = detector(image)
+        
+        if len(faces) != 2:
+            st.error("Se necesitan exactamente dos caras para realizar el intercambio.")
+            return None
+
+        landmarks1 = predictor(image, faces[0])
+        landmarks2 = predictor(image, faces[1])
+        points1 = np.array([[p.x, p.y] for p in landmarks1.parts()])
+        points2 = np.array([[p.x, p.y] for p in landmarks2.parts()])
+
+        mask1 = np.zeros(image.shape[:2], dtype=np.uint8)
+        cv2.fillConvexPoly(mask1, cv2.convexHull(points1), 255)
+        mask2 = np.zeros(image.shape[:2], dtype=np.uint8)
+        cv2.fillConvexPoly(mask2, cv2.convexHull(points2), 255)
+        face1 = cv2.bitwise_and(image, image, mask=mask1)
+        face2 = cv2.bitwise_and(image, image, mask=mask2)
+        image[mask1 == 255] = face2[mask1 == 255]
+        image[mask2 == 255] = face1[mask2 == 255]
+        swapped_image_path = DATA_PATH / f"swapped_image_{Path(image_path).stem}.jpg"
+        cv2.imwrite(str(swapped_image_path), image)
+        return str(swapped_image_path)
     except Exception as e:
         st.error(f"Error en el face swap: {e}")
         return None
@@ -161,54 +177,35 @@ def main():
     steps = st.sidebar.slider("Pasos", 1, 100, 20)
     seed = st.sidebar.number_input("Semilla", value=-1)
 
-    if format_option == "9:16":
-        width = 720
-        height = 1280
-    else:
-        width = 1280
-        height = 720
+    width, height = (720, 1280) if format_option == "9:16" else (1280, 720)
 
     if st.sidebar.button("Generar Imagen"):
         with st.spinner("Mejorando y generando imagen..."):
             result = asyncio.run(gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language))
-            image_paths = result[0]
-            prompt_file = result[1]
+            image_paths, prompt_file = result[0], result[1]
 
         st.write(f"Image paths: {image_paths}")
 
+        if image_paths and Path(image_paths).exists():
+            st.image(image_paths, caption="Imagen generada", use_column_width=True)
+        if prompt_file and Path(prompt_file).exists():
+            with open(prompt_file, "r") as file:
+                st.text_area("Prompt utilizado", file.read(), height=150)
+
+    st.sidebar.header("Galería de Imágenes")
+    image_storage, usage = get_storage()
+    st.sidebar.write(usage)
+    for img_path in image_storage:
+        st.sidebar.image(img_path, width=100)
+    
+    if st.sidebar.button("Borrar Imagen"):
+        delete_image(image_paths)
+
+    if st.sidebar.button("Intercambiar Caras"):
         if image_paths:
-            if Path(image_paths).exists():
-                st.image(image_paths, caption="Imagen Generada")
-            else:
-                st.error("El archivo de imagen no existe.")
-
-            if prompt_file and Path(prompt_file).exists():
-                prompt_text = Path(prompt_file).read_text()
-                st.write(f"Prompt utilizado: {prompt_text}")
-            else:
-                st.write("El archivo del prompt no está disponible.")
-
-    files, usage = get_storage()
-    st.text(usage)
-    cols = st.columns(6)
-    prompts = get_prompts()
-
-    for idx, file in enumerate(files):
-        with cols[idx % 6]:
-            image = Image.open(file)
-            prompt_file = prompts.get(Path(file).stem.replace("image_", ""), None)
-            prompt_text = Path(prompt_file).read_text() if prompt_file else "No disponible"
-            
-            st.image(image, caption=f"Imagen {idx+1}")
-            st.write(f"Prompt: {prompt_text}")
-            
-            if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx+1}"):
-                delete_image(file)
-            
-            if st.button(f"Swap Face {idx+1}", key=f"swap_{idx+1}"):
-                swapped_image_path = swap_faces(file)
-                if swapped_image_path:
-                    st.image(swapped_image_path, caption=f"Imagen {idx+1} con Face Swap")
+            swapped_path = swap_faces(image_paths)
+            if swapped_path:
+                st.image(swapped_path, caption="Imagen con caras intercambiadas", use_column_width=True)
 
 if __name__ == "__main__":
     main()