Update app.py

app.py

@@ -8,6 +8,7 @@ from huggingface_hub import HfApi, create_repo
 # 環境変数から取得
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_REPO_USER = os.environ.get("MODEL_REPO", None)  # ユーザー名
+LORA_REPO = os.environ.get("LORA_REPO", None)  # LoRA用リポジトリ
 CIVITAI_TOKEN = os.environ.get("CIVITAI_TOKEN", None)  # Civitai用トークン
 
 api = HfApi()
@@ -16,41 +17,33 @@ def download_and_upload(model_id, repo_type, model_name=None):
     if not model_id or not repo_type or (repo_type == "model" and not model_name):
         return "モデルIDまたはタイプ、モデル名が未指定です。"
 
-    # ①クローン元リポジトリの定義
-    source_repo = "John6666/pony-realism-v21main-sdxl"
-    print(f"クローン元リポジトリ: {source_repo}")
-
-    # ②新規リポジトリの定義
-    target_repo = f"{MODEL_REPO_USER}/{model_name}"
-    print(f"新規リポジトリ: {target_repo}")
-
-    # 新規リポジトリを作成
-    try:
-        print(f"新規リポジトリを作成します: {target_repo}")
-        create_repo(repo_id=target_repo, token=HF_TOKEN, private=False)
-        print(f"新規リポジトリを作成しました: {target_repo}")
-    except Exception as e:
-        print(f"新規リポジトリの作成に失敗しました: {e}")
-        return f"新規リポジトリの作成に失敗しました: {e}"
-
-    # クローン元リポジトリをクローン
-    try:
-        print(f"クローン元リポジトリをクローン中: {source_repo}")
-        subprocess.run(["git", "clone", f"https://huggingface.co/{source_repo}"], check=True)
-        print(f"クローン元リポジトリをクローンしました: {source_repo}")
-    except Exception as e:
-        print(f"クローン元リポジトリのクローンに失敗しました: {e}")
-        return f"クローン元リポジトリのクローンに失敗しました: {e}"
-
-    # クローンしたリポジトリをコピー先として使用
-    source_repo_name = source_repo.split("/")[-1]
-    os.chdir(source_repo_name)
-    print(f"現在の作業ディレクトリ: {os.getcwd()}")
-
-    # CIVITAIからファイルをダウンロード
+    # LoRAの場合
+    if repo_type == "lora":
+        if not LORA_REPO:
+            return "LoRA用リポジトリが環境変数で設定されていません。"
+        target_repo = LORA_REPO
+        print(f"LoRAファイルのアップロード先リポジトリ: {target_repo}")
+
+    # モデルの場合
+    else:
+        # 新規リポジトリを定義
+        target_repo = f"{MODEL_REPO_USER}/{model_name}"
+        print(f"モデルファイルのアップロード先リポジトリ: {target_repo}")
+
+        # 新規リポジトリを作成
+        try:
+            print(f"新規リポジトリを作成します: {target_repo}")
+            create_repo(repo_id=target_repo, token=HF_TOKEN, private=False)
+            print(f"新規リポジトリを作成しました: {target_repo}")
+        except Exception as e:
+            print(f"新規リポジトリの作成に失敗しました: {e}")
+            return f"新規リポジトリの作成に失敗しました: {e}"
+
+    # Civitaiからファイルをダウンロード
     headers = {}
     if CIVITAI_TOKEN:
         headers["Authorization"] = f"Bearer {CIVITAI_TOKEN}"
+
     try:
         response = requests.get(f"https://civitai.com/api/v1/models/{model_id}", headers=headers)
         response.raise_for_status()
@@ -74,38 +67,55 @@ def download_and_upload(model_id, repo_type, model_name=None):
         print(f"Civitaiモデルファイルのダウンロードに失敗しました: {e}")
         return f"Civitaiモデルファイルのダウンロードに失敗しました: {e}"
 
-    # ファイルを移動してGit操作
+    # LoRAまたはモデルリポジトリにファイルをアップロード
     try:
-        subprocess.run(["git", "add", "."], check=True)
-        subprocess.run(["git", "commit", "-m", f"Add {filename}"], check=True)
-        subprocess.run(["git", "push"], check=True)
-        print(f"ファイルを新規リポジトリ '{target_repo}' にアップロードしました。")
-        return f"ファイルを新規リポジトリ '{target_repo}' にアップロードしました。"
+        # ファイルアップロード
+        api.upload_file(
+            path_or_fileobj=filename,
+            path_in_repo=filename,
+            repo_id=target_repo,
+            token=HF_TOKEN,
+            repo_type="model" if repo_type == "model" else "dataset",
+            commit_message=f"Add {filename} from model ID {model_id}"
+        )
+        print(f"ファイルをアップロードしました: {filename}")
+        return f"ファイル '{filename}' をリポジトリ '{target_repo}' にアップロードしました。"
     except Exception as e:
-        print(f"Git操作中にエラーが発生しました: {e}")
-        return f"Git操作中にエラーが発生しました: {e}"
+        print(f"アップロード中���エラーが発生しました: {e}")
+        return f"アップロード中にエラーが発生しました: {e}"
 
 # Gradio UI構築
 with gr.Blocks() as demo:
     gr.Markdown("## CivitaiモデルIDからファイルダウンロード＆アップロードツール")
-    gr.Markdown("CivitaiのモデルIDを指定して、新規モデルリポジトリを作成します。")
+    gr.Markdown("CivitaiのモデルIDを指定して、新規モデルまたはLoRAリポジトリへアップロードします。")
 
     model_id_input = gr.Textbox(
         label="CivitaiモデルID",
         placeholder="例: 918104"
     )
     repo_choice = gr.Radio(
-        choices=["model"],
+        choices=["model", "lora"],
         label="アップロード先タイプの選択",
         value="model"
     )
     model_name_input = gr.Textbox(
-        label="新規リポジトリ名",
-        placeholder="例: my-awesome-model"
+        label="新規モデルリポジトリ名（modelを選択した場合に必須）",
+        placeholder="例: my-awesome-model",
+        visible=False
     )
     run_button = gr.Button("実行")
     output = gr.Textbox(label="結果メッセージ", interactive=False)
 
+    # モデルが選択された場合にモデル名入力を表示
+    def toggle_model_name(repo_type):
+        return gr.update(visible=(repo_type == "model"))
+
+    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, model_name_input],