Update app.py

app.py

@@ -1,90 +1,196 @@
 import os
-import re
-import requests
 import gradio as gr
-from huggingface_hub import HfApi
+import requests
+import re
+
+from huggingface_hub import HfApi, list_repo_files
+import datasets  # pip install datasets
+from datasets import Dataset
 
-# 環境変数から取得
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL_REPO = os.environ.get("MODEL_REPO", None)
+MODEL_REPO = os.environ.get("MODEL_REPO", None)  # => "username/my_dataset_repo"
 LORA_REPO = os.environ.get("LORA_REPO", None)
-CIVITAI_TOKEN = os.environ.get("CIVITAI_TOKEN", None)  # Civitai用トークン
+CIVITAI_TOKEN = os.environ.get("CIVITAI_TOKEN", None)
 
 api = HfApi()
 
+#################################
+# ファイルアップロード部分 (既存)
+#################################
 def download_and_upload(url, repo_type):
-    # 入力チェック
-    if not url or not repo_type:
-        return "URLまたはタイプが未指定です。"
+    """
+    既存コード。Civitai APIからファイルをダウンロードし、
+    Datasetリポジトリのfiles/以下にアップロードするイメージ。
+    """
+    # (詳しくは省略。前回のコードを参照)
+
+#################################
+# メタデータ操作部分
+#################################
 
-    # Civitai API用ヘッダ
-    headers = {}
-    if CIVITAI_TOKEN:
-        headers["Authorization"] = f"Bearer {CIVITAI_TOKEN}"
+def load_metadata(repo_type):
+    """
+    ロードボタンで呼ばれ、Datasetリポジトリ内のmetadataを取得してUIへ表示。
+    metadataが存在しない場合は作成する。
+    """
+    if repo_type == "model":
+        target_repo = MODEL_REPO
+    else:
+        target_repo = LORA_REPO
+
+    if not target_repo:
+        return None, f"{repo_type}リポジトリが設定されていません。"
 
-    # ファイルダウンロード実行
-    # CivitaiのAPIからモデルを取得
     try:
-        response = requests.get(url, headers=headers)
-        response.raise_for_status()  # ステータスコードがエラーなら例外
+        # list_repo_filesでリポジトリにmetadata.arrowがあるか確認
+        files_in_repo = list_repo_files(repo_id=target_repo, repo_type="dataset", token=HF_TOKEN)
+        has_metadata = any("metadata/" in f for f in files_in_repo)
+
+        if has_metadata:
+            # 既存のmetadataデータセットをロード
+            ds = datasets.load_dataset(
+                path=target_repo,        # "username/dataset_name"
+                data_dir="metadata",     # metadata/下にあるarrowデータを読み込む
+                use_auth_token=HF_TOKEN,
+                repo_type="dataset"
+            )
+            # ds はDictionary形式( split名 -> Dataset )を返す。例えば ds["train"] が本体。
+            dset = ds["train"]
+            # Dataframe表示用に2次元リストに変換
+            # カラム: "filename", "order" など
+            df_data = []
+            for row in dset:
+                df_data.append([row["filename"], row.get("order", 0)])
+            return df_data, f"メタデータをロードしました（{len(dset)}ファイル）"
+        else:
+            # metadataが無い → 初期化する
+            # すでにfiles/下にファイルがあるかもしれないため、それらをリストアップ
+            all_files = [f for f in files_in_repo if not f.startswith(".")]
+
+            # filesディレクトリに限る
+            file_list = [f for f in all_files if f.startswith("files/")]
+
+            # 初期メタデータ作成
+            # "filename" 列と "order" 列を持つDatasetにする
+            metadata_dict = {
+                "filename": file_list,
+                "order": list(range(len(file_list)))
+            }
+            dset = Dataset.from_dict(metadata_dict)
+            # push_to_hub で metadata/ 以下に格納
+            dset.push_to_hub(
+                repo_id=target_repo,
+                token=HF_TOKEN,
+                split="train",
+                repo_type="dataset",
+                private=False,
+                embed_remote_files=False,
+                max_shard_size="50MB",  # メタデータなので小さいはず
+                dataset_type=None,      # arrowデフォルト
+                data_dir="metadata",    # ここがmetadataディレクトリ
+                # push_to_hub は datasets>=2.11.0 以降推奨
+            )
+            # UI表示用の2次元リスト
+            df_data = [[fn, i] for fn, i in zip(file_list, range(len(file_list)))]
+            return df_data, "メタデータを新規作成しました。"
     except Exception as e:
-        return f"ファイルのダウンロードに失敗しました: {e}"
-
-    # Content-Dispositionからfilenameを抽出
-    # 例: Content-Disposition: attachment; filename="model.safetensors"
-    content_disp = response.headers.get("content-disposition", "")
-    filename = None
-    if content_disp:
-        # filename="..." という���分を正規表現で抽出
-        match = re.search(r'filename="?([^"]+)"?', content_disp)
-        if match:
-            filename = match.group(1)
-
-    if not filename:
-        # ファイル名がヘッダに無い場合のfallback
-        filename = "downloaded_file.bin"
-    
-    # ダウンロードファイルをローカル一時ファイルに保存
-    with open(filename, "wb") as f:
-        f.write(response.content)
-
-    # 選択されたリポジトリを決定
+        return None, f"ロード失敗: {e}"
+
+def apply_reorder(table_data, repo_type):
+    """
+    UIで編集したDataframe (filename, order) を受け取り、
+    Datasetを再構築してpush_to_hub。ファイルの物理名は変えず、メタデータだけ更新。
+    """
+    if not table_data or len(table_data) == 0:
+        return "テーブルが空です。"
+
     if repo_type == "model":
         target_repo = MODEL_REPO
     else:
         target_repo = LORA_REPO
 
     if not target_repo:
-        return f"{repo_type}用リポジトリが環境変数で設定されていません。"
+        return f"{repo_type}リポジトリが設定されていません。"
+
+    # table_data: [[filename1, order1], [filename2, order2], ...]
+    # Datasetに再変換
+    filenames = []
+    orders = []
+    for row in table_data:
+        if len(row) < 2:
+            continue
+        filenames.append(row[0])
+        # orderをintに変換
+        try:
+            orders.append(int(row[1]))
+        except:
+            orders.append(0)
+
+    # ここでファイル名自体は変えず、並び順(order)のカラムだけ変える
+    metadata_dict = {"filename": filenames, "order": orders}
+    dset = Dataset.from_dict(metadata_dict)
 
-    # Hugging Face Hub にアップロード
     try:
-        api.upload_file(
-            path_or_fileobj=filename,
-            path_in_repo=filename,
+        # push_to_hubで上書きコミット
+        dset.push_to_hub(
             repo_id=target_repo,
             token=HF_TOKEN,
-            repo_type="dataset",  # dataset リポジトリとして扱う
-            commit_message=f"Add {filename} from {url}"
+            split="train",
+            repo_type="dataset",
+            private=False,
+            embed_remote_files=False,
+            data_dir="metadata",   # metadataディレクトリに保存
+            commit_message="Update metadata order"
         )
-        return f"ファイル '{filename}' を '{target_repo}' にアップロードしました。"
+        return f"メタデータを更新しました（{len(dset)}ファイル）。"
     except Exception as e:
-        return f"アップロード中にエラーが発生しました: {e}"
-    finally:
-        # 必要ならローカルファイル削除などの後処理
-        pass
+        return f"並び替えのコミットに失敗: {e}"
 
-# Gradio UI構築 (前回とほぼ同様)
+##########################
+# Gradio UI定義
+##########################
 with gr.Blocks() as demo:
-    gr.Markdown("## Civitai経由ファイルダウンロード＆アップロードツール")
-    gr.Markdown("Civitai APIのURLを指定して、'model'または'lora'リポジトリへアップロードします。")
+    gr.Markdown("# データセットリポジトリ管理（ファイル名そのままで並び順編集）")
+
+    with gr.Row():
+        url_input = gr.Textbox(label="CivitaiモデルAPIのURL", placeholder="例: https://civitai.com/api/download/models/xxx?type=Model&format=SafeTensor")
+        repo_choice = gr.Radio(choices=["model", "lora"], label="アップロード先タイプ", value="model")
+    upload_button = gr.Button("アップロード実行")
+    upload_result = gr.Textbox(label="アップロード結果")
+
+    # 既存のアップロード処理を呼ぶ
+    upload_button.click(
+        download_and_upload,
+        inputs=[url_input, repo_choice],
+        outputs=upload_result
+    )
+
+    gr.Markdown("## メタデータによる並び替え管理")
+
+    load_button = gr.Button("ロード（メタデータ）")
+    file_table = gr.Dataframe(
+        headers=["filename", "order"],
+        row_count=0,
+        col_count=2,
+        datatype=["str", "number"],
+        wrap=True,
+        label="ファイル + 並び順メタデータ"
+    )
+    load_result = gr.Textbox(label="ロード結果")
+
+    load_button.click(
+        fn=load_metadata,
+        inputs=[repo_choice],
+        outputs=[file_table, load_result]
+    )
 
-    url_input = gr.Textbox(label="CivitaiモデルAPIのURL", placeholder="例: https://civitai.com/api/download/models/xxxxx?type=Model&format=SafeTensor")
-    repo_choice = gr.Radio(choices=["model", "lora"], label="アップロード先タイプの選択", value="model")
-    run_button = gr.Button("実行")
-    output = gr.Textbox(label="結果メッセージ", interactive=False)
+    reorder_button = gr.Button("並び順を反映")
+    reorder_result = gr.Textbox(label="並び替え実行結果")
 
-    run_button.click(download_and_upload, inputs=[url_input, repo_choice], outputs=output)
+    reorder_button.click(
+        fn=apply_reorder,
+        inputs=[file_table, repo_choice],
+        outputs=reorder_result
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.queue()