Hoaihao/drive-zero-gpu-middleware

1import gradio as gr2from googleapiclient.discovery import build3from google.oauth2 import service_account4from googleapiclient.http import MediaIoBaseDownload5from gradio_client import Client, handle_file6import io7import os8import time9import datetime10import requests11import base6412import threading 13import random14import json15from concurrent.futures import ThreadPoolExecutor16 17# --- CẤU HÌNH TỪ KÉT SẮT (SECRETS) ---18INPUT_FOLDER_ID = os.getenv("INPUT_ID")19HISTORY_FOLDER_ID = os.getenv("HISTORY_ID")20WEB_APP_URL = os.getenv("WEB_URL")21 22# --- KẾT NỐI ĐẾN APP CỘNG ĐỒNG ---23CLIENT_URL = "sczhou/CodeFormer" 24 25# BIẾN TOÀN CỤC LƯU LOG26SYSTEM_LOGS = ">>> Hệ thống V12: Cấu hình chuẩn Studio (Fidelity 0.7 - Giảm ảo)...\n"27 28# --- HÀM KẾT NỐI DRIVE (AN TOÀN) ---29def get_drive_service_safe():30    try:31        json_str = os.getenv("GDRIVE_JSON")32        if not json_str: return None33        info = json.loads(json_str)34        creds = service_account.Credentials.from_service_account_info(info, scopes=['https://www.googleapis.com/auth/drive'])35        return build('drive', 'v3', credentials=creds)36    except Exception as e:37        print(f"Lỗi Connect Drive: {e}")38        return None39 40# --- WORKER XỬ LÝ 1 ẢNH ---41def process_single_image(file_item):42    global SYSTEM_LOGS43    44    file_id = file_item['id']45    file_name = file_item['name']46    47    drive_service = get_drive_service_safe()48    if not drive_service: return49 50    thread_id = threading.get_ident()51    input_path = f"temp_{thread_id}_{file_name}"52    output_path = None53 54    try:55        SYSTEM_LOGS += f"-> [BẮT ĐẦU] Gửi '{file_name}' đi xử lý...\n"56        57        # 1. Tải ảnh58        request = drive_service.files().get_media(fileId=file_id)59        fh = io.BytesIO()60        downloader = MediaIoBaseDownload(fh, request)61        done = False62        while done is False: status, done = downloader.next_chunk()63        64        with open(input_path, "wb") as f: f.write(fh.getbuffer())65        66        # 2. Gửi sang App Cộng Đồng67        # Sử dụng Client ẩn danh (không token)68        client = Client(CLIENT_URL)69        70        # --- CẤU HÌNH "CHUYÊN GIA" ĐỂ GIẢM ẢO ---71        # Tham số 1: Ảnh72        # Tham số 2 (Face Align): True (Căn chỉnh mặt cho thẳng)73        # Tham số 3 (Background Enhance): True (Làm nét cả nền)74        # Tham số 4 (Face Upsample): True (Làm nét chi tiết mặt)75        # Tham số 5 (Upscale): 2 (Phóng to 2 lần - Đủ dùng cho in ấn)76        # Tham số 6 (Fidelity): 0.7 (QUAN TRỌNG NHẤT: 0 là ảo, 1 là thật. 0.7 là đẹp nhất)77        78        result = client.predict(79            handle_file(input_path), 80            True,                    81            True,                    82            True,                    83            2,                       84            0.7,                     # <--- CHỈNH SỐ NÀY ĐỂ HẾT ẢO (Tăng lên 0.8 hoặc 0.9 nếu vẫn thấy ảo)85            api_name="/predict"86        )87        88        if result:89            output_path = result[0] if isinstance(result, (list, tuple)) else result90        else:91            raise ValueError("Không nhận được ảnh trả về")92 93        # 3. Gửi về Web App94        with open(output_path, "rb") as image_file:95            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')96        97        final_output_name = f"Done_{file_name}"98        99        payload = {100            "filename": final_output_name,101            "file": encoded_string,102            "mimeType": "image/png",103            "is_result": True 104        }105        106        requests.post(WEB_APP_URL, json=payload)107        SYSTEM_LOGS += f"-> [XONG] Đã trả ảnh: {final_output_name}\n"108        109        # 4. Lưu kho110        random_suffix = random.randint(100000, 999999)111        timestamp = int(time.time())112        if "Anhpng.com" in file_name:113            new_name = f"Processed_{timestamp}_{file_name}"114        else:115            ext = file_name.split('.')[-1]116            new_name = f"Anhpng.com - {timestamp}_{random_suffix}.{ext}"117        118        drive_service.files().update(119            fileId=file_id,120            addParents=HISTORY_FOLDER_ID,121            removeParents=INPUT_FOLDER_ID,122            body={'name': new_name}123        ).execute()124 125    except Exception as e:126        err_msg = str(e)127        if "queue" in err_msg.lower():128            SYSTEM_LOGS += f"!!! {file_name}: App đang bận xếp hàng, thử lại sau.\n"129        else:130            SYSTEM_LOGS += f"!!! Lỗi xử lý {file_name}: {err_msg[:100]}...\n"131    132    finally:133        if os.path.exists(input_path): os.remove(input_path)134        try:135            if output_path and os.path.exists(output_path): os.remove(output_path)136        except: pass137 138# --- QUẢN LÝ CHÍNH ---139def main_manager():140    global SYSTEM_LOGS141    142    # Giữ 2 luồng để ổn định143    MAX_WORKERS = 2 144    145    while True:146        time.sleep(1) 147        if len(SYSTEM_LOGS) > 20000:148            SYSTEM_LOGS = "--- [Auto Clean Log] ---\n" + SYSTEM_LOGS[-10000:]149 150        try:151            service = get_drive_service_safe()152            if not service: continue153 154            try:155                results = service.files().list(156                    q=f"'{INPUT_FOLDER_ID}' in parents and trashed=false",157                    fields="files(id, name)").execute()158                items = results.get('files', [])159            except:160                continue161 162            if not items: continue163 164            with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:165                executor.map(process_single_image, items)166                167        except Exception as e:168            print(f"Lỗi Manager: {str(e)}")169            time.sleep(5)170 171# --- KÍCH HOẠT ---172threading.Thread(target=main_manager, daemon=True).start()173 174# --- GIAO DIỆN ---175def get_logs():176    return SYSTEM_LOGS[-2000:]177 178with gr.Blocks() as demo:179    gr.Markdown("### Middleware V12 (Fidelity 0.7 - Realistic Mode)")180    status = gr.Textbox(label="Log Trạng thái", lines=15)181    timer = gr.Timer(1) 182    timer.tick(get_logs, outputs=status)183 184demo.launch()