Update main.py

main.py

@@ -1,122 +1,127 @@
-from quart import Quart, request, jsonify, send_from_directory, websocket
-from gradio_client import Client, file
-from quart_cors import cors
-import os
-import traceback
 import shutil
 import base64
-import asyncio
-import httpx
-from tenacity import retry, stop_after_attempt, wait_fixed
+from gradio_client import Client, file
+
+
+
+
 
-app = Quart(__name__)
-cors(app)
 
+app = Flask(__name__)
+CORS(app)
 # 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)
+UPLOAD_FOLDER = tempfile.mkdtemp()
+RESULT_FOLDER = tempfile.mkdtemp()
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['RESULT_FOLDER'] = RESULT_FOLDER
 
-@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
-async def initialize_client():
-    return Client("yisol/IDM-VTON")
 
-client = None
 
-@app.before_serving
-async def startup():
-    global client
-    try:
-        client = await initialize_client()
-        print("Client initialized successfully")
-    except Exception as e:
-        print(f"Failed to initialize client: {str(e)}")
+
+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),
+                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 "Prediction timed out after {} seconds".format(timeout)
+
+    if isinstance(result[0], Exception):
+        return str(result[0])  # Return the error message
+    return result[0]
+
+@app.route('/')
+def index():
+
+    return {'message': 'This is a wearon API'}
 
 @app.route('/process', methods=['POST'])
-async def predict():
-    global client
-    if client is None:
-        return jsonify(error='Client not initialized. Please try again later.'), 503
+def predict():
 
     try:
-        form = await request.form
-        files = await request.files
-
         # Get the product image URL from the request
-        product_image_url = form.get('product_image_url')
+        product_image_url = request.form.get('product_image_url')
+        if not product_image_url:
+
+            return jsonify(error='No product image URL provided'), 400
         
         # Handle the uploaded model image
-        if 'model_image' not in files:
+        if 'model_image' not in request.files:
+
             return jsonify(error='No model image file provided'), 400
 
-        model_image = 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))
+        full_filename = os.path.abspath(filename)
 
-        print("Product image = ", product_image_url)
-        print("Model image = ", full_filename)
+        print("Product image URL:", product_image_url)
+        print("Model image path:", 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
+        # Perform prediction with a timeout
+        result = predict_with_timeout(full_filename, product_image_url)
+        if isinstance(result, str):
+
+            return jsonify(error=result), 500
+
+        print("Prediction result:", result)
 
-        print(result)
-        # Extract the path of the first output image
+        # Check if the result contains a valid path
         output_image_path = result[0]
+        if not os.path.exists(output_image_path):
+            return jsonify(error='Output image file not found: {}'.format(output_image_path)), 500
+
 
         # 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)
 
+
         # Remove the uploaded file after processing
-        os.remove(filename)
+        os.remove(full_filename)
+
 
         # 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
 
 @app.route('/uploads/<filename>')
-async def uploaded_file(filename):
-    return await send_from_directory(app.config['UPLOAD_FOLDER'], filename)
+def uploaded_file(filename):
 
-@app.websocket('/ws')
-async def ws():
-    while True:
-        data = await websocket.receive()
-        await websocket.send(f"Echo: {data}")
+    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+    app.run(host='0.0.0.0', port=7860)