Update vton.py

vton.py

@@ -2,6 +2,10 @@ import replicate
 import requests
 import os
 import base64
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
 
 def virtual_try_on(garm_img_data, human_img_data, garment_des):
     """
@@ -34,24 +38,41 @@ def virtual_try_on(garm_img_data, human_img_data, garment_des):
         input=input_data
     )
     
-    # The model returns a list of URLs, get the first one
-    if not output or not isinstance(output, list):
-        raise ValueError("Invalid response from the model")
-        
-    output_url = output[0]
-    
-    # Validate URL scheme
-    if not output_url.startswith(('http://', 'https://')):
-        output_url = 'https://' + output_url
-    
-    # Download the image from the URL
-    response = requests.get(output_url)
-    response.raise_for_status()  # Raise an exception for bad status codes
+    # Validate and process the model output
+    if not output:
+        raise ValueError("Empty response received from the model")
     
     # Save the output image
     output_path = "output.jpg"
-    with open(output_path, "wb") as file:
-        file.write(response.content)
+    
+    # Handle FileOutput type from Replicate
+    if hasattr(output, 'save'):
+        # If output is a FileOutput object, save it directly
+        output.save(output_path)
+    else:
+        # Handle URL-based response (fallback to existing logic)
+        if isinstance(output, str):
+            output_url = output
+        elif isinstance(output, list) and output:
+            output_url = output[0]
+        else:
+            # Instead of raising an error, try to get the content directly
+            try:
+                with open(output_path, "wb") as file:
+                    file.write(output.read())
+                return output_path
+            except Exception as e:
+                raise ValueError(f"Unable to process model output: {str(e)}")
+            
+        # Validate URL scheme
+        if not output_url.startswith(('http://', 'https://')) and not output_url.startswith('http'):
+            output_url = 'https://' + output_url.lstrip('/')
+        # Download the image from the URL
+        response = requests.get(output_url)
+        response.raise_for_status()  # Raise an exception for bad status codes
+        
+        with open(output_path, "wb") as file:
+            file.write(response.content)
     
     print(f"Virtual try-on complete! Output saved to {output_path}")
     return output_path