Update app.py

app.py

@@ -3,18 +3,18 @@ import json
 import re
 import requests
 import gradio as gr
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi, hf_hub_download, RepositoryNotFoundError
 
 # 環境変数から取得
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL_REPO = os.environ.get("MODEL_REPO", None)
+MODEL_REPO_USERNAME = os.environ.get("MODEL_REPO", None)  # ユーザー名を想定
 LORA_REPO = os.environ.get("LORA_REPO", None)
 CIVITAI_TOKEN = os.environ.get("CIVITAI_TOKEN", None)  # Civitai用トークン
 
 api = HfApi()
 
-def download_and_upload(model_id, repo_type):
-    print(f"モデルID: {model_id}, リポジトリタイプ: {repo_type}")
+def download_and_upload(model_id, repo_type, model_name=None):
+    print(f"モデルID: {model_id}, リポジトリタイプ: {repo_type}, モデル名: {model_name}")
     if not model_id or not repo_type:
         return "モデルIDまたはタイプが未指定です。"
 
@@ -64,17 +64,42 @@ def download_and_upload(model_id, repo_type):
     print(f"ファイルを保存しました: {filename}")
 
     if repo_type == "model":
-        target_repo = MODEL_REPO
+        if not MODEL_REPO_USERNAME:
+            return "MODEL_REPO環境変数にユーザー名が設定されていません。"
+        if not model_name:
+            return "モデル名が未指定です。"
+
+        target_repo = f"{MODEL_REPO_USERNAME}/{model_name}"
+        repo_api_type = "model"
     else:
         target_repo = LORA_REPO
+        repo_api_type = "dataset"
 
     if not target_repo:
         return f"{repo_type}用リポジトリが環境変数で設定されていません。"
 
-    # リポジトリ内のファイル一覧を取得
+    # リポジトリが存在するか確認、存在しない場合は作成
     try:
-        repo_files = api.list_repo_files(repo_id=target_repo, repo_type="dataset", token=HF_TOKEN)
-        print(f"リポジトリ内のファイル一覧を取得しました。")
+        repo_files = api.list_repo_files(repo_id=target_repo, repo_type=repo_api_type, token=HF_TOKEN)
+        print(f"リポジトリ '{target_repo}' が存在します。")
+    except RepositoryNotFoundError:
+        if repo_type == "model":
+            try:
+                api.create_repo(
+                    name=model_name,
+                    token=HF_TOKEN,
+                    repo_type=repo_api_type,
+                    exist_ok=True
+                )
+                print(f"リポジトリ '{target_repo}' を新規作成しました。")
+                repo_files = []
+            except Exception as e:
+                print(f"リポジトリの作成に失敗しました: {e}")
+                return f"リポジトリの作成に失敗しました: {e}"
+        else:
+            print(f"リポジトリ '{target_repo}' が存在しません。")
+            return f"リポジトリ '{target_repo}' が存在しません。"
+
     except Exception as e:
         print(f"リポジトリ内のファイル一覧の取得に失敗しました: {e}")
         return f"リポジトリ内のファイル一覧の取得に失敗しました: {e}"
@@ -87,17 +112,24 @@ def download_and_upload(model_id, repo_type):
         metadata_file = hf_hub_download(
             repo_id=target_repo,
             filename=metadata_filename,
-            repo_type="dataset",
+            repo_type=repo_api_type,
             token=HF_TOKEN,
             force_download=True
         )
         with open(metadata_file, 'r', encoding='utf-8') as f:
             metadata = json.load(f)
         print("既存のメタデータを読み込みました。")
-    except Exception as e:
+    except RepositoryNotFoundError:
+        # リポジトリ自体が存在しない場合は既に処理済み
+        metadata = []
+        print("メタデータファイルが存在しないため、新規作成します。")
+    except FileNotFoundError:
         # メタデータファイルが存在しない場合は新規作成
         metadata = []
         print("メタデータファイルが存在しないため、新規作成します。")
+    except Exception as e:
+        print(f"メタデータの読み込みに失敗しました: {e}")
+        return f"メタデータの読み込みに失敗しました: {e}"
 
     # メタデータを更新（存在しないファイルを削除）
     updated_metadata = []
@@ -142,7 +174,7 @@ def download_and_upload(model_id, repo_type):
             path_in_repo=filename,
             repo_id=target_repo,
             token=HF_TOKEN,
-            repo_type="dataset",
+            repo_type=repo_api_type,
             commit_message=f"Add {filename} from model ID {model_id}"
         )
 
@@ -152,7 +184,7 @@ def download_and_upload(model_id, repo_type):
             path_in_repo=metadata_filename,
             repo_id=target_repo,
             token=HF_TOKEN,
-            repo_type="dataset",
+            repo_type=repo_api_type,
             commit_message="Update metadata"
         )
         print(f"ファイルとメタデータをアップロードしました。")
@@ -166,22 +198,47 @@ with gr.Blocks() as demo:
     gr.Markdown("## CivitaiモデルIDからファイルダウンロード＆アップロードツール")
     gr.Markdown("CivitaiのモデルIDを指定して、'model'または'lora'リポジトリへアップロードします。")
 
-    # ダウンロード＆アップロード機能
-    model_id_input = gr.Textbox(
-        label="CivitaiモデルID",
-        placeholder="例: 1102"
-    )
-    repo_choice = gr.Radio(
-        choices=["model", "lora"],
-        label="アップロード先タイプの選択",
-        value="model"
-    )
-    run_button = gr.Button("実行")
-    output = gr.Textbox(label="結果メッセージ", interactive=False)
+    with gr.Row():
+        model_id_input = gr.Textbox(
+            label="CivitaiモデルID",
+            placeholder="例: 1102",
+            lines=1
+        )
+
+    with gr.Row():
+        repo_choice = gr.Radio(
+            choices=["model", "lora"],
+            label="アップロード先タイプの選択",
+            value="model",
+            interactive=True
+        )
+
+    with gr.Row():
+        model_name_input = gr.Textbox(
+            label="モデル名",
+            placeholder="例: my-awesome-model",
+            visible=True,
+            lines=1
+        )
+
+    with gr.Row():
+        run_button = gr.Button("実行")
+    
+    with gr.Row():
+        output = gr.Textbox(label="結果メッセージ", interactive=False, lines=2)
+
+    # リポジトリタイプの選択に応じてモデル名入力を表示/非表示
+    def toggle_model_name(repo_type):
+        if repo_type == "model":
+            return gr.update(visible=True)
+        else:
+            return gr.update(visible=False)
+
+    repo_choice.change(toggle_model_name, inputs=repo_choice, outputs=model_name_input)
 
     run_button.click(
         download_and_upload,
-        inputs=[model_id_input, repo_choice],
+        inputs=[model_id_input, repo_choice, model_name_input],
         outputs=output
     )