diff --git a/.gitignore b/.gitignore index 3b72af06f63e14ea65f6aac0929f38aa57590d94..372c1ae0de201a12fec788f8840e6c3b09a164ad 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ .env __pycache__ +service_account.json diff --git a/requirements.txt b/requirements.txt index 10a5b34a0239f84647ba66c386fd134a08d38847..8d02e887dc0efaff2f454a1a2d86369938b4bac7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,8 @@ +dotenv>=0.9.9 Flask>=2.0.3 flask-cors>=5.0.1 -dotenv>=0.9.9 +google-api-python-client>=2.161.0 +google-auth>=2.38.0 gradio_client>=1.7.0 pydantic>=1.8.2 requests>=2.26.0 diff --git a/src/controllers/gdrive_controller.py b/src/controllers/gdrive_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..2a111ca101d01400a0a4b31b1284b6f9598090bc --- /dev/null +++ b/src/controllers/gdrive_controller.py @@ -0,0 +1,34 @@ +import os +from flask import jsonify, Response +from googleapiclient.discovery import build +from googleapiclient.http import MediaFileUpload +from google.oauth2 import service_account + +class GDriveController: + def __init__(self): + SERVICE_ACCOUNT_FILE = "service_account.json" + SCOPES = ["https://www.googleapis.com/auth/drive.file"] + + creds = service_account.Credentials.from_service_account_file(SERVICE_ACCOUNT_FILE, scopes=SCOPES) + self.drive_service = build("drive", "v3", credentials=creds) + + def post_gdrive(self, local_file_path: str, folder_id: str) -> Response: + if not local_file_path or not folder_id: + return jsonify({"error": "Missing local_file_path or folder_id"}), 400 + + if not os.path.exists(local_file_path): + return jsonify({"error": "File not found"}), 404 + + file_metadata = { + "name": os.path.basename(local_file_path), + "parents": [folder_id] + } + media = MediaFileUpload(local_file_path, resumable=True) + + file = self.drive_service.files().create( + body=file_metadata, + media_body=media, + fields="id, webViewLink" + ).execute() + + return jsonify({"file_url": file.get('webViewLink')}) diff --git a/src/routes/router.py b/src/routes/router.py index bb778bf622d211a80b845ef59c63ed2d3e9f837e..314864ede51b25c2f9820142038bcc4ce160e0aa 100644 --- a/src/routes/router.py +++ b/src/routes/router.py @@ -1,12 +1,15 @@ -from flask import Blueprint, request +from flask import Blueprint, Response, request +from controllers.gdrive_controller import GDriveController from controllers.glb_controller import GlbController from controllers.render_controller import RenderController from controllers.suggest_controller import SuggestController +gdrive_bp = Blueprint('gdrive', __name__) glb_bp = Blueprint('glb', __name__) render_bp = Blueprint('render', __name__) suggest_bp = Blueprint('suggest', __name__) +gdrive_controller = GDriveController() glb_controller = GlbController() render_controller = RenderController() suggest_controller = SuggestController() @@ -14,7 +17,7 @@ suggest_controller = SuggestController() def setup_routes(app): @app.route('/glb', methods=['GET']) # Image to model - def get_glb(): + def get_glb() -> Response: hf_token = request.args.get('hf_token', '', type=str) base_image_url = request.args.get('base_image_url', type=str) ss_guidance_strength = request.args.get('ss_guidance_strength', 7.5, type=float) @@ -30,9 +33,9 @@ def setup_routes(app): slat_sampling_steps, multiimage_algo, mesh_simplify, texture_size) - # Text to image + # Text + Sketch to image @app.route('/render', methods=['GET']) - def get_render(): + def get_render() -> Response: hf_token = request.args.get('hf_token', '', type=str) base_image_url = request.args.get('base_image_url', '', type=str) prompt = request.args.get('prompt', 'A realistic photo of a building', type=str) @@ -46,7 +49,7 @@ def setup_routes(app): # Suggest @app.route('/suggest', methods=['GET']) - def get_suggest(): + def get_suggest() -> Response: hf_token = request.args.get('hf_token', '', type=str) message = request.args.get('message', 'Hello!!', type=str) system_prompt = request.args.get('system_prompt', 'You are a helpful assistant. First recognize the user request and then reply carefully with thinking.', type=str) @@ -60,3 +63,11 @@ def setup_routes(app): temperature, max_new_tokens, top_k_sampling, repetition_penalty, top_p_sampling) + + # Upload to Google Drive + @app.route('/upload', methods=['POST']) + def post_gdrive() -> Response: + local_file_path = request.form.get('local_file_path', '', type=str) + folder_id = request.form.get('folder_id', '', type=str) + + return gdrive_controller.post_gdrive(local_file_path, folder_id)