app.py

import gradio as gr
import os
import zipfile
import time
from datetime import datetime
from huggingface_hub import login, hf_hub_download, HfApi

# --- Config ---
REPO_ID = os.environ.get("HFPATH")
HF_TOKEN = os.environ.get("MAGIC")

# Đăng nhập vào Hugging Face
login(HF_TOKEN)

# Khởi tạo API
api = HfApi()

# --- Hàm xóa file JPG ---
def delete_jpg_files():
    print("Chạy lệnh xóa")
    # Lấy danh sách file trong repo
    files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
    
    # Lọc file JPG
    jpg_files = [file for file in files if file.endswith('.jpg') and '/' not in file]
    #print("Chờ 1 giờ trước khi tiếp tục...")
    #time.sleep(3600)  # Chờ 1 giờ
    #api.delete_file(repo_id=REPO_ID, path_in_repo=jpg_files, repo_type="dataset")
    #print(f"Đã xóa file: {jpg_files}")  # In ra console
    # Xóa file theo từng lô
    for i in range(0, len(jpg_files), 99):
        batch = jpg_files[i:i + 99]  # Lấy 99 file một lần
        for file in batch:
            api.delete_file(repo_id=REPO_ID, path_in_repo=file, repo_type="dataset")
            print(f"Đã xóa file: {file}")  # In ra console
        if i + 99 < len(jpg_files):  # Nếu còn file để xóa
            print("Chờ 1 giờ trước khi tiếp tục...")
            time.sleep(3600)  # Chờ 1 giờ
    return "Xóa thành công."

# --- Hàm tải file từ Hugging Face ---
def download_from_hf(subfolder):
    downloaded_files = []
    
    # Lấy danh sách tất cả file trong repo
    all_files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
    
    # Lọc file theo subfolder
    for file in all_files:
        if file.startswith(subfolder + "/"):
            # Tải file
            downloaded_file = hf_hub_download(repo_id=REPO_ID, filename=file, repo_type="dataset")
            downloaded_files.append(downloaded_file)
    
    return downloaded_files

# --- Hàm tạo file ZIP ---
def create_zip(files):
    zip_filename = f"downloaded_files_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for file in files:
            zipf.write(file, os.path.basename(file))
    return zip_filename

# --- Hàm chính cho Gradio ---
def download_and_zip(subfolder):
    #delete_jpg_files()  # Xóa file JPG trước khi tải
    files = download_from_hf(subfolder)  # Tải file từ subfolder
    if files:
        zip_file = create_zip(files)  # Tạo file ZIP
        return zip_file
    else:
        return "Không tìm thấy file trong subfolder đã chỉ định."

# --- Giao diện Gradio ---
iface = gr.Interface(
    fn=download_and_zip,
    inputs=gr.Textbox(label="Tên Subfolder"),
    outputs=gr.File(label="Tải File Zip"),
    title="Trình Tải File Hugging Face",
    description="Nhập tên subfolder để tải file từ dataset Hugging Face và tạo file zip."
)

iface.launch()