Update app.py

app.py

@@ -1,11 +1,10 @@
 # ===================================================================
 #
-#   【最終版v3】AIペルソナ選択式 Gradioアプリ (環境自動判定対応)
+#   【最終完成版v4】AIペルソナ選択式 Gradioアプリ
 #
 # ===================================================================
 
 # ライブラリのインポート
-print("⚙️ 必要なライブラリをインポートします...")
 import torch
 import torch.nn as nn
 import warnings
@@ -18,59 +17,52 @@ from peft import PeftModel
 import gradio as gr
 from PIL import Image
 
-# ログレベルを設定
+# --- グローバル変数と基本設定 ---
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
 warnings.filterwarnings("ignore")
 
-# --- グローバル変数 ---
 MODELS_CACHE = {}
 CURRENT_PERSONA = None
-BASE_MODEL = None  # ベースモデルをグローバルに保持
+BASE_MODEL = None
 TOKENIZER = None
 CLIP_MODEL = None
 CLIP_PROCESSOR = None
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
-# ★★★ 実行環境を自動で判定し、モデルへのパスを正しく設定します ★★★
-# ★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
+# --- 環境判定とパス設定 ---
 if 'SPACE_ID' in os.environ:
-    print("✅ Hugging Face Spaces環境で実行中")
-    # Hugging Face Spacesでは、アップロードされたファイルはルートにあります
+    logger.info("✅ Hugging Face Spaces環境で実行中")
     MODEL_BASE_PATH = '.'
 else:
     try:
         from google.colab import drive
-        print("\n🔗 Googleドライブをマウントします...")
+        logger.info("\n🔗 Googleドライブをマウントします...")
         drive.mount('/content/drive', force_remount=True)
         MODEL_BASE_PATH = '/content/drive/MyDrive'
-        print("✅ Google Colab環境で実行中")
+        logger.info("✅ Google Colab環境で実行中")
     except ImportError:
-        print("⚠️ ローカル環境として実行します。'./models'フォルダにモデルを配置してください。")
-        # ローカルPCで実行する場合、このコードと同じ階層に'models'フォルダを作成し、
-        # その中に`final_model_`で始まるフォルダ群を配置してください。
+        logger.warning("⚠️ ローカル環境として実行します。'./models'フォルダにモデルを配置してください。")
         MODEL_BASE_PATH = './models'
 
-print(f"🖥️ 使用デバイス: {DEVICE}")
-print(f"📂 モデルデータの検索パス: {MODEL_BASE_PATH}")
+logger.info(f"🖥️ 使用デバイス: {DEVICE}")
+logger.info(f"📂 モデルデータの検索パス: {MODEL_BASE_PATH}")
 
 
 # --- モデル定義 ---
 class MultimodalModel(nn.Module):
-    def __init__(self, base_model, face_embedding_dim):
+    def __init__(self, peft_model, face_embedding_dim):
         super().__init__()
-        self.base_model = base_model
-        self.face_injector = nn.Linear(face_embedding_dim, self.base_model.config.hidden_size)
+        self.peft_model = peft_model
+        self.face_injector = nn.Linear(face_embedding_dim, self.peft_model.config.hidden_size)
 
     def forward(self, input_ids, attention_mask, face_embedding, **kwargs):
         try:
-            target_device = self.base_model.get_input_embeddings().weight.device
-            target_dtype = self.base_model.get_input_embeddings().weight.dtype
+            target_device = self.peft_model.device
+            target_dtype = self.peft_model.get_input_embeddings().weight.dtype
 
             self.face_injector.to(target_device, dtype=target_dtype)
-            inputs_embeds = self.base_model.get_input_embeddings()(input_ids)
+            inputs_embeds = self.peft_model.get_input_embeddings()(input_ids)
             injected_face_features = self.face_injector(face_embedding.to(target_device, dtype=target_dtype))
 
             conditioned_embeds = torch.cat([injected_face_features.unsqueeze(1), inputs_embeds], dim=1)
@@ -78,11 +70,8 @@ class MultimodalModel(nn.Module):
                 torch.ones(attention_mask.shape[0], 1, device=target_device, dtype=attention_mask.dtype),
                 attention_mask
             ], dim=1)
-
-            if "inputs" in kwargs:
-                 kwargs["inputs"] = kwargs["inputs"].to(target_device)
-
-            return self.base_model.generate(
+            
+            return self.peft_model.generate(
                 inputs_embeds=conditioned_embeds,
                 attention_mask=new_attention_mask,
                 **kwargs
@@ -92,46 +81,16 @@ class MultimodalModel(nn.Module):
             traceback.print_exc()
             raise
 
-# --- モデルロード関連の関数 ---
-@torch.no_grad()
-def get_available_models():
-    """利用可能なモデルのリストを取得"""
-    try:
-        if not os.path.exists(MODEL_BASE_PATH):
-            logger.warning(f"モデルディレクトリが存在しません: {MODEL_BASE_PATH}")
-            return []
-
-        model_paths = glob.glob(os.path.join(MODEL_BASE_PATH, "final_model_*"))
-        if not model_paths:
-            logger.warning(f"モデルファイルが見つかりません: {MODEL_BASE_PATH}に final_model_* がありません。")
-            return []
-
-        model_names = sorted([os.path.basename(p).replace("final_model_", "").replace("_", " ") for p in model_paths])
-        return model_names
-    except Exception as e:
-        logger.error(f"モデル検索エラー: {e}")
-        return []
+# ★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
+# ★★★ エラー修正箇所 ★★★
+# ★★★ モデルロードの仕組みを全面的に改善しました ★★★
+# ★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
 
 @torch.no_grad()
-def load_model_and_dependencies(person_name, progress=gr.Progress()):
-    """全ての依存モデルと指定された人物モデルをロード"""
-    global BASE_MODEL, FINAL_MODEL, CURRENT_PERSONA, TOKENIZER, CLIP_MODEL, CLIP_PROCESSOR
-
+def initialize_core_models():
+    """アプリ起動時に一度だけ、重いコアモデルをロードする"""
+    global BASE_MODEL, TOKENIZER, CLIP_MODEL, CLIP_PROCESSOR
     try:
-        # --- 共通モデル（Tokenizer, CLIP, Base LLM）の初期化 ---
-        progress(0, desc="🔄 コアモデルを準備中...")
-        if TOKENIZER is None:
-            logger.info("📝 Tokenizerをロード中...")
-            TOKENIZER = AutoTokenizer.from_pretrained("rinna/nekomata-7b-instruction", trust_remote_code=True)
-            TOKENIZER.pad_token = TOKENIZER.eos_token
-            logger.info("✅ Tokenizerのロード完了")
-
-        if CLIP_MODEL is None or CLIP_PROCESSOR is None:
-            logger.info("📷 CLIP画像エンコーダーをロード中...")
-            CLIP_MODEL = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(DEVICE)
-            CLIP_PROCESSOR = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
-            logger.info("✅ CLIPモデルのロード完了")
-        
         if BASE_MODEL is None:
             logger.info("🧠 ベースLLMをロード中... (時間がかかります)")
             quantization_config = BitsAndBytesConfig(
@@ -145,8 +104,29 @@ def load_model_and_dependencies(person_name, progress=gr.Progress()):
             )
             logger.info("✅ ベースLLMのロード完了")
 
-        # --- ペルソナモデルのロード ---
-        progress(0.5, desc=f"🔄 AI人格「{person_name}」を適用中...")
+        if TOKENIZER is None:
+            logger.info("📝 Tokenizerをロード中...")
+            TOKENIZER = AutoTokenizer.from_pretrained("rinna/nekomata-7b-instruction", trust_remote_code=True)
+            TOKENIZER.pad_token = TOKENIZER.eos_token
+            logger.info("✅ Tokenizerのロード完了")
+
+        if CLIP_MODEL is None or CLIP_PROCESSOR is None:
+            logger.info("📷 CLIP画像エンコーダーをロード中...")
+            CLIP_MODEL = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(DEVICE)
+            CLIP_PROCESSOR = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
+            logger.info("✅ CLIPモデルのロード完了")
+        return True
+    except Exception as e:
+        logger.error(f"コアモデルの初期化エラー: {e}")
+        traceback.print_exc()
+        return False
+
+@torch.no_grad()
+def load_persona_model(person_name, progress):
+    """指定されたAI人格（アダプター）をロードする"""
+    global FINAL_MODEL, CURRENT_PERSONA
+    try:
+        progress(0.1, desc=f"🔄 AI人格「{person_name}」を準備中...")
         if person_name == CURRENT_PERSONA and FINAL_MODEL is not None:
             logger.info(f"✅ AI人格「{person_name}」は既に準備完了です。")
             return True
@@ -185,7 +165,7 @@ def load_model_and_dependencies(person_name, progress=gr.Progress()):
         return True
 
     except Exception as e:
-        logger.error(f"モデルロードおよび依存関係の初期化エラー: {e}")
+        logger.error(f"AI人格のロードエラー: {e}")
         traceback.print_exc()
         return False
 
@@ -197,7 +177,7 @@ def predict(person_name, image, instruction, max_len, temp, top_p, progress=gr.P
         if image is None: return "❌ エラー: 画像をアップロードしてください。"
         if not instruction.strip(): return "❌ エラー: 指示（プロンプト）を入力してください。"
 
-        if not load_model_and_dependencies(person_name, progress):
+        if not load_persona_model(person_name, progress):
             return f"❌ エラー: AI人格「{person_name}」のロードに失敗しました。詳細はログを確認してくださ���。"
         
         progress(0.4, desc="🖼️ 顔の雰囲気を分析中...")
@@ -306,9 +286,9 @@ if __name__ == "__main__":
             logger.info("🌐 Gradioアプリケーションを起動します...")
             demo.launch(debug=False) 
     else:
-        # 初期化に失敗した場合、エラーUIを表示する
+        logger.error("❌ コアモデルの初期化に失敗したため、アプリを起動できません。")
+        # エラーが発生した場合でも、簡単なエラー表示UIを起動する
         with gr.Blocks(title="起動エラー") as error_demo:
              gr.Markdown("# ❌ 致命的なエラー: コアモデルの初期化に失敗しました。")
-             gr.Markdown("ログを確認して、必要なライブラリがインストールされているか、モデルへのパスが正しいかを確認してください。")
-        error_demo.launch(debug=False)
-        logger.error("❌ コアモデルの初期化に失敗したため、アプリを起動できません。")
+             gr.Markdown("アプリケーションのログを確認して、エラーの原因を特定してください。")
+        error_demo.launch(debug=False)