Update app.py

app.py

@@ -1,70 +1,46 @@
 import gradio as gr
+import requests
 from PIL import Image
-from diffusers import AutoPipelineForInpainting, AutoencoderKL
-import torch
-from SegBody import segment_body  # Import the segmentation function
-
-# Check if CUDA is available and set the device accordingly
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load models with the correct precision based on the device
-if device == "cuda":
-    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)  # Use fp16 for GPU
-    pipeline = AutoPipelineForInpainting.from_pretrained(
-        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
-        vae=vae, 
-        torch_dtype=torch.float16,  # Use fp16 for GPU
-        variant="fp16",  # Correct variant for GPU
-        use_safetensors=True
-    ).to(device)  # Ensure it uses the GPU
-else:
-    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32)  # Use fp32 for CPU
-    pipeline = AutoPipelineForInpainting.from_pretrained(
-        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
-        vae=vae, 
-        torch_dtype=torch.float32,  # Use fp32 for CPU
-        variant="fp16",  # Use fp32 for CPU
-        use_safetensors=True
-    ).to(device)  # Ensure it uses the CPU if no GPU
-
-# Define the inference function
-def inpaint(person_image, garment_image, prompt):
-    # Preprocess the images by resizing them to 512x512
-    person_image = person_image.convert("RGB").resize((512, 512))
-    garment_image = garment_image.convert("RGB").resize((512, 512))
-
-    # Use segment_body to generate the body mask for inpainting
-    seg_image, mask_image = segment_body(person_image, face=False)  # You can control face removal here (face=False)
-
-    # Resize mask to 512x512 to match the inpainting requirements
-    mask_image = mask_image.resize((512, 512))
-
-    # Perform inpainting using the pipeline
-    results = pipeline(
-        prompt=prompt,
-        negative_prompt="ugly, bad quality, bad anatomy",
-        image=person_image,
-        mask_image=mask_image,  # Use the mask from segmentation
-        ip_adapter_image=garment_image,  # Garment image as the IP Adapter image
-        strength=0.99,
-        guidance_scale=8.0,
-        num_inference_steps=100
-    )
-
-    return results.images[0]  # Return the generated image
-
-# Set up the Gradio interface
-demo = gr.Interface(
-    fn=inpaint,
+import io
+
+# URL del modelo TryOnGAN (o cualquier otro modelo Hugging Face)
+MODEL_URL = "https://huggingface.co/spaces/akhaliq/TryOnGAN"
+
+# Función que procesa la imagen usando la API de Hugging Face
+def try_on_clothes(photo, clothing):
+    # Convertir las imágenes a formato adecuado para la solicitud
+    files = {
+        'photo': ('photo.jpg', photo, 'image/jpeg'),
+        'clothing': ('clothing.jpg', clothing, 'image/jpeg')
+    }
+    
+    try:
+        # Enviar las imágenes al modelo de Hugging Face
+        response = requests.post(MODEL_URL, files=files)
+        
+        if response.status_code == 200:
+            # Suponiendo que la respuesta es una imagen generada (puede variar según el modelo)
+            output_image = Image.open(io.BytesIO(response.content))
+            return output_image
+        else:
+            return "Error al procesar las imágenes, intente nuevamente."
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Crear la interfaz de Gradio
+iface = gr.Interface(
+    fn=try_on_clothes,
     inputs=[
-        gr.Image(type="pil", label="Person Image"),  # Input for person image
-        gr.Image(type="pil", label="Garment Image"),  # Input for garment image
-        gr.Textbox(label="Prompt", placeholder="Enter the prompt for the model")  # Text prompt for inpainting
+        gr.Image(type="pil", label="Foto del Maniquí"),
+        gr.Image(type="pil", label="Prenda a Probar")
     ],
-    outputs=gr.Image(type="pil"),
-    title="Stable Diffusion Inpainting with Segmentation",
-    description="Inpainting model for seamless garment transfer on segmented body image using Stable Diffusion XL.",
-    server_timeout=100,  # Increase timeout duration to prevent session errors
+    outputs=gr.Image(type="pil", label="Resultado del Probador Virtual"),
+    title="Probador Virtual AI",
+    description="Sube una foto de tu maniquí y una prenda para probarla en tiempo real con TryOnGAN."
 )
 
-demo.launch(share=True)  # Enable share link for testing in a public domain
+# Ejecutar la interfaz
+if __name__ == "__main__":
+    iface.launch()
+
+