Update src/main_MM.py

src/main_MM.py

@@ -11,17 +11,22 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI
 import os
 from flask import Flask, request, jsonify
 from flask_swagger_ui import get_swaggerui_blueprint
+
+from routing.injectionRouter import injection_router
+from routing.detoxifyRouter import detoxify_router
+from routing.embedingRouter import embed_router
+from routing.privacyRouter import privacy_router
+from routing.topicRouter import topic_router
+from routing.sentimentRouter import sentiment_router
+from routing.invisibletextRouter import invisibletext_router
+from routing.gibberishRouter import gibberish_router
+from routing.bancodeRouter import bancode_router
+# from routing.safety_router import img_router
 from config.logger import CustomLogger, request_id_var
 from waitress import serve
-from werkzeug.exceptions import HTTPException, UnsupportedMediaType
+from werkzeug.exceptions import HTTPException, UnsupportedMediaType, BadRequest
 from mapper.mapper import *
 
-# ----------------------
-# Create necessary log directories to avoid runtime warnings
-os.makedirs("/responsible-ai/logs", exist_ok=True)
-os.makedirs("/aicloud/logs", exist_ok=True)
-# ----------------------
-
 SWAGGER_URL = '/rai/v1/raimoderationmodels/docs'  
 API_URL = '/static/swagger.json'
 
@@ -30,87 +35,61 @@ swaggerui_blueprint = get_swaggerui_blueprint(
     API_URL
 )
 
+# ------------------- LOG DIRECTORY FIX -------------------
+# Use /tmp which is writable in Hugging Face Spaces
+LOG_DIR_RESPONSIBLE_AI = os.path.join("/tmp", "responsible-ai", "logs")
+LOG_DIR_AICLOUD = os.path.join("/tmp", "aicloud", "logs")
+os.makedirs(LOG_DIR_RESPONSIBLE_AI, exist_ok=True)
+os.makedirs(LOG_DIR_AICLOUD, exist_ok=True)
+# Optional: pass these directories to your logger if needed
 log = CustomLogger()
+# ----------------------------------------------------------
+
 app = Flask(__name__)
 app.register_blueprint(swaggerui_blueprint)
+app.register_blueprint(injection_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(detoxify_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(embed_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(privacy_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(topic_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(sentiment_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(invisibletext_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(gibberish_router, url_prefix='/rai/v1/raimoderationmodels')
+app.register_blueprint(bancode_router, url_prefix='/rai/v1/raimoderationmodels')
+# app.register_blueprint(img_router, url_prefix='/rai/v1/raimoderationmodels')
+
 
-# ----------------------
-# Function to load all models before starting server
-def load_models():
-    from routing.injectionRouter import init_injection_model
-    from routing.detoxifyRouter import init_detoxify_model
-    from routing.embedingRouter import init_embed_model
-    from routing.privacyRouter import init_privacy_model
-    from routing.topicRouter import init_topic_model
-
-    init_injection_model()
-    init_detoxify_model()
-    init_embed_model()
-    init_privacy_model()
-    init_topic_model()
-# ----------------------
-
-# ----------------------
-# Function to register all routers
-def register_routers():
-    from routing.injectionRouter import injection_router
-    from routing.detoxifyRouter import detoxify_router
-    from routing.embedingRouter import embed_router
-    from routing.privacyRouter import privacy_router
-    from routing.topicRouter import topic_router
-    from routing.sentimentRouter import sentiment_router
-    from routing.invisibletextRouter import invisibletext_router
-    from routing.gibberishRouter import gibberish_router
-    from routing.bancodeRouter import bancode_router
-    # from routing.safety_router import img_router
-
-    routers = [
-        injection_router,
-        detoxify_router,
-        embed_router,
-        privacy_router,
-        topic_router,
-        sentiment_router,
-        invisibletext_router,
-        gibberish_router,
-        bancode_router
-        # img_router
-    ]
-
-    for r in routers:
-        app.register_blueprint(r, url_prefix='/rai/v1/raimoderationmodels')
-# ----------------------
-
-# ----------------------
-# Error handling
 def handle_http_exception(exc):
+    """Handles HTTP exceptions, returning a JSON response."""
     response = jsonify({"error": exc.description})
     response.status_code = exc.code
     return response
 
 @app.errorhandler(HTTPException)
 def handle_all_http_exceptions(exc):
+    """Global exception handler for HTTP errors."""
     return handle_http_exception(exc)
 
 def handle_unsupported_mediatype(exc):
-    return jsonify({"error": "Unsupported media type"}), 415
+    """Handles unsupported media type exceptions."""
+    return jsonify({"error": "Unsupported media type"}), 415  # 415 Unsupported Media Type
 
 @app.errorhandler(UnsupportedMediaType)
 def handle_all_unsupported_mediatype_exceptions(exc):
+    """Global exception handler for unsupported media types."""
     return handle_unsupported_mediatype(exc)
-# ----------------------
 
-if __name__ == "__main__":
-    load_models()       # Load models before registering routes
-    register_routers()  # Register routers after models are loaded
 
+if __name__ == "__main__":
     serve(
-        app,
-        host='0.0.0.0',
-        port=7860,
+        app, 
+        host='0.0.0.0', 
+        port=7860, 
         threads=int(os.getenv('THREADS', 1)),
-        connection_limit=int(os.getenv('CONNECTION_LIMIT', 500)),
+        connection_limit=int(os.getenv('CONNECTION_LIMIT', 500)), 
         channel_timeout=int(os.getenv('CHANNEL_TIMEOUT', 120))
     )
+    # app.run()
+