Update main.py

main.py

@@ -1,26 +1,70 @@
 import tempfile
-import os
-from flask import Flask, request, jsonify
+import threading
+import time
+import signal
+from flask import Flask, request, jsonify, send_from_directory
 from flask_cors import CORS
-from diffusers import DiffusionPipeline
+import os
+import traceback
+import shutil
+import base64
+from gradio_client import Client, file
+
 
 app = Flask(__name__)
 CORS(app)
 
-# Load the diffusion model
-pipeline = DiffusionPipeline.from_pretrained("yisol/IDM-VTON")
+# Define Gradio client instance
+client = Client("yisol/IDM-VTON")
 
-# Directory to save uploaded files
+# Directory to save uploaded and processed files
 UPLOAD_FOLDER = tempfile.mkdtemp()
+RESULT_FOLDER = tempfile.mkdtemp()
+if not os.path.exists(UPLOAD_FOLDER):
+    os.makedirs(UPLOAD_FOLDER)
+if not os.path.exists(RESULT_FOLDER):
+    os.makedirs(RESULT_FOLDER)
+
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+app.config['RESULT_FOLDER'] = RESULT_FOLDER
+
+def predict_with_timeout(model_image_path, product_image_url, timeout=600):
+    result = [None]  # Mutable object to store the result
+
+    def target():
+        try:
+            result[0] = client.predict(
+                dict({"background": file(model_image_path), "layers": [], "composite": None}),
+                garm_img=file(product_image_url),
+                garment_des="Hello!!",
+                is_checked=True,
+                is_checked_crop=False,
+                denoise_steps=30,
+                seed=42,
+                api_name="/tryon"
+            )
+        except Exception as e:
+            result[0] = str(e)
+
+    thread = threading.Thread(target=target)
+    thread.start()
+    thread.join(timeout)
+    if thread.is_alive():
+        return None  # Timeout
+    if isinstance(result[0], Exception):
+        return str(result[0])  # Return the error message
+    return result[0]
 
 @app.route('/')
 def index():
-    return {'hello': 'This is a wearon API using Diffusers'}
+    return {'hello': 'This is a wearon API'}
 
 @app.route('/process', methods=['POST'])
-def process():
+def predict():
     try:
+        # Get the product image URL from the request
+        product_image_url = request.form.get('product_image_url')
+        
         # Handle the uploaded model image
         if 'model_image' not in request.files:
             return jsonify(error='No model image file provided'), 400
@@ -33,24 +77,40 @@ def process():
         filename = os.path.join(app.config['UPLOAD_FOLDER'], model_image.filename)
         model_image.save(filename)
 
-        # Get the product image URL from the request
-        product_image_url = request.form.get('product_image_url')
+        base_path = os.getcwd()
+        full_filename = os.path.normpath(os.path.join(base_path, filename))
+
+        print("Product image = ", product_image_url)
+        print("Model image = ", full_filename)
+        
+        # Perform prediction with a timeout
+        result = predict_with_timeout(full_filename, product_image_url)
+        if result is None:
+            return jsonify(error='Prediction timed out after 10 minutes'), 500
+
+        print(result)
+        # Extract the path of the first output image
+        output_image_path = result[0]
 
-        # Here you would process the images with your pipeline
-        # Example (you would replace this with the actual method to process the images):
-        result = pipeline(product_image_url, model_image_path=filename)
+        # Copy the output image to the RESULT_FOLDER
+        output_image_filename = os.path.basename(output_image_path)
+        local_output_path = os.path.join(app.config['RESULT_FOLDER'], output_image_filename)
+        shutil.copy(output_image_path, local_output_path)
 
-        # Assuming the pipeline returns a path to the result image
-        result_image_path = result["path"]
+        # Remove the uploaded file after processing
+        os.remove(filename)
 
-        # Read the result image and encode it in base64
-        with open(result_image_path, "rb") as image_file:
+        # Encode the output image in base64
+        with open(local_output_path, "rb") as image_file:
             encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
 
+        # Return the output image in JSON format
         return jsonify(image=encoded_image), 200
 
     except Exception as e:
+        traceback.print_exc()
         return jsonify(error=str(e)), 500
 
-if __name__ == '__main__':
-    app.run(debug=True)
+@app.route('/uploads/<filename>')
+def uploaded_file(filename):
+    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)