Update main.py

main.py

@@ -1,100 +1,71 @@
-from quart import Quart, request, jsonify, send_from_directory
-from quart_cors import cors
+from flask import Flask, request, jsonify, send_from_directory
 from gradio_client import Client, file
+from flask_cors import CORS
 import os
 import traceback
 import shutil
 import base64
 import asyncio
-import aiofiles
 
-app = Quart(__name__)
-cors(app)
+app = Flask(__name__)
+CORS(app)
 
 client = Client("kadirnar/IDM-VTON")
 
-# Directory to save uploaded and processed files
 UPLOAD_FOLDER = 'static/uploads'
 RESULT_FOLDER = 'static/results'
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-if not os.path.exists(RESULT_FOLDER):
-    os.makedirs(RESULT_FOLDER)
-
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(RESULT_FOLDER, exist_ok=True)
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['RESULT_FOLDER'] = RESULT_FOLDER
 
-@app.route('/')
-async def home():
-    return {"wearon": "wearon model is running!"}
-
 @app.route('/process', methods=['POST'])
 async def predict():
     try:
-        # Get the product image URL from the request
-        form = await request.form
-        product_image_url = form.get('product_image_url')
-
-        # Handle the uploaded model image
+        product_image_url = request.form.get('product_image_url')
+        
         if 'model_image' not in request.files:
             return jsonify(error='No model image file provided'), 400
-
-        model_image = await request.files['model_image']
+        
+        model_image = request.files['model_image']
         if model_image.filename == '':
             return jsonify(error='No selected file'), 400
-
-        # Save the uploaded file to the upload directory
+        
         filename = os.path.join(app.config['UPLOAD_FOLDER'], model_image.filename)
-        await model_image.save(filename)
-
+        model_image.save(filename)
+        
         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
-        try:
-            result = await asyncio.to_thread(client.predict,
-                                             dict={"background": file(full_filename), "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:
-            traceback.print_exc()
-            raise
-
+        
+        print("Product image = ", product_image_url)
+        print("Model image = ", full_filename)
+        
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(None, client.predict, {
+            "background": file(full_filename), "layers": [], "composite": None
+        }, file(product_image_url), "Hello!!", True, False, 30, 42, "/tryon")
+        
         print(result)
-        # Extract the path of the first output image
         output_image_path = result[0]
-
-        # 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)
-        await asyncio.to_thread(shutil.copy, output_image_path, local_output_path)
-
-        # Remove the uploaded file after processing
+        shutil.copy(output_image_path, local_output_path)
+        
         os.remove(filename)
-
-        # Encode the output image in base64
-        async with aiofiles.open(local_output_path, "rb") as image_file:
-            encoded_image = base64.b64encode(await image_file.read()).decode('utf-8')
-
-        # Return the output image in JSON format
+        
+        with open(local_output_path, "rb") as image_file:
+            encoded_image = base64.b64encode(image_file.read()).decode('utf-8')
+        
         return jsonify(image=encoded_image), 200
-
+    
     except Exception as e:
         traceback.print_exc()
         return jsonify(error=str(e)), 500
 
 @app.route('/uploads/<filename>')
-async def uploaded_file(filename):
-    return await send_from_directory(app.config['UPLOAD_FOLDER'], filename)
+def uploaded_file(filename):
+    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)