Fix #22

app.py

@@ -10,6 +10,7 @@ import random
 from torchvision.utils import save_image
 import gradio as gr
 import numpy as np
+import io
 
 # Asegúrate de que las funciones necesarias estén definidas (si no lo están ya)
 def resize(img, size):
@@ -251,7 +252,7 @@ class Solver(object):
             print(f"Error al cargar el checkpoint: {e}.")
             raise Exception(f"Error al cargar el checkpoint: {e}")
 
-    def transfer_style(self, source_image, reference_image): # Eliminado target_domain_index
+    def transfer_style(self, source_image, reference_image):
         # Asegúrate de que los modelos estén en modo de evaluación
         self.G.eval()
         self.S.eval()
@@ -264,17 +265,14 @@ class Solver(object):
                 transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
             ])
             # Convertir a PIL image antes de la transformación
-            source_image = Image.fromarray(source_image)
-            reference_image = Image.fromarray(reference_image)
+            source_image = Image.open(io.BytesIO(source_image))  # Use BytesIO
+            reference_image = Image.open(io.BytesIO(reference_image))
 
             source_image = transform(source_image).unsqueeze(0).to(self.device)
             reference_image = transform(reference_image).unsqueeze(0).to(self.device)
 
-            # Crear el tensor de dominio objetivo
-            # target_domain = torch.tensor([target_domain_index]).to(self.device) # Eliminado
-
             # Codificar el estilo de la imagen de referencia
-            s_ref = self.S(reference_image, torch.tensor([0]).to(self.device)) # Simplificado
+            s_ref = self.S(reference_image, torch.tensor([0]).to(self.device))
 
             # Generar la imagen con el estilo transferido
             generated_image = self.G(source_image, s_ref)
@@ -284,7 +282,7 @@ class Solver(object):
             return generated_image
 
 # Función principal para la inferencia
-def main(source_image, reference_image, checkpoint_path, args): # Eliminado target_domain_index
+def main(source_image, reference_image, checkpoint_path, args):
     if source_image is None or reference_image is None:
         raise gr.Error("Por favor, proporciona ambas imágenes (fuente y referencia).")
 
@@ -294,37 +292,35 @@ def main(source_image, reference_image, checkpoint_path, args): # Eliminado targ
     solver.load_checkpoint(checkpoint_path)
 
     # Realizar la transferencia de estilo
-    generated_image = solver.transfer_style(source_image, reference_image) # Eliminado target_domain_index
+    generated_image = solver.transfer_style(source_image, reference_image)
     return generated_image
 
 def gradio_interface():
     # Definir los argumentos (ajustados para la inferencia)
     args = SimpleNamespace(
-        img_size=128,  # Asegúrate de que esto coincida con el tamaño de imagen usado en el entrenamiento
-        num_domains=3, # Cambiado a 3 para que coincida con el checkpoint del MappingNetwork
-        latent_dim=16, # Puedes ajustar esto si es necesario
+        img_size=128,
+        num_domains=3,
+        latent_dim=16,
         style_dim=64,
-        num_workers=0,  # Establecer en 0 para evitar problemas en algunos entornos
+        num_workers=0,
         seed=8365,
     )
 
     # Ruta al checkpoint
-    checkpoint_path = "iter/10500_nets_ema.ckpt"  # Reemplaza con la ruta correcta a tu checkpoint
+    checkpoint_path = "iter/20500_nets_ema.ckpt"
 
     # Crear la interfaz de Gradio
     inputs = [
         gr.Image(label="Source Image (Car to change style)"),
         gr.Image(label="Reference Image (Style to transfer)"),
-        # gr.Radio(choices=[0, 1, 2], label="Target Domain (0: BMW, 1: Corvette, 2: Mazda)", value=0), # Eliminado
     ]
     outputs = gr.Image(label="Generated Image (Car with transferred style)")
 
     title = "AutoStyleGAN: Car Style Transfer"
-    description = "Transfer the style of one car to another. Upload a source car image and a reference car image." # Modificado
+    description = "Transfer the style of one car to another. Upload a source car image and a reference car image."
 
-    # Crear la interfaz de Gradio
     iface = gr.Interface(
-        fn=lambda source_image, reference_image: main(source_image, reference_image, checkpoint_path, args), # Eliminado target_domain_index
+        fn=lambda source_image, reference_image: main(source_image, reference_image, checkpoint_path, args),
         inputs=inputs,
         outputs=outputs,
         title=title,
@@ -334,4 +330,4 @@ def gradio_interface():
 
 if __name__ == '__main__':
     iface = gradio_interface()
-    iface.launch(share=True)
+    iface.launch(share=True)