.

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 from huggingface_hub import AsyncInferenceClient
 from typing import List, Dict, Optional, Union
 import logging
-from dataclasses import dataclass
 from enum import Enum, auto
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, pipeline
@@ -15,397 +14,366 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# モデルの型定義
-class ModelType(Enum):
-    LOCAL = "local"
-    INFERENCE_API = "inference_api"
+# モデルタイプの定義
+LOCAL = "local"
+INFERENCE_API = "inference_api"
 
-@dataclass
-class ModelConfig:
-    name: str
-    description: str
-    type: ModelType
-    model_id: Optional[str] = None
-    model_path: Optional[str] = None
-
-# モデル定義を拡充
+# モデル定義
 TEXT_GENERATION_MODELS = [
-    ModelConfig(
-        name="Zephyr-7B",
-        description="Specialized in understanding context and nuance",
-        type=ModelType.INFERENCE_API,
-        model_id="HuggingFaceH4/zephyr-7b-beta"
-    ),
-    ModelConfig(
-        name="Llama-2",
-        description="Known for its robust performance in content analysis",
-        type=ModelType.LOCAL,
-        model_path="meta-llama/Llama-2-7b-hf"
-    ),
-    ModelConfig(
-        name="Mistral-7B",
-        description="Offers precise and detailed text evaluation",
-        type=ModelType.LOCAL,
-        model_path="mistralai/Mistral-7B-v0.1"
-    )
+    {
+        "name": "Zephyr-7B",
+        "description": "Specialized in understanding context and nuance",
+        "type": INFERENCE_API,
+        "model_id": "HuggingFaceH4/zephyr-7b-beta"
+    },
+    {
+        "name": "Llama-2",
+        "description": "Known for its robust performance in content analysis",
+        "type": LOCAL,
+        "model_path": "meta-llama/Llama-2-7b-hf"
+    },
+    {
+        "name": "Mistral-7B",
+        "description": "Offers precise and detailed text evaluation",
+        "type": LOCAL,
+        "model_path": "mistralai/Mistral-7B-v0.1"
+    }
 ]
 
 CLASSIFICATION_MODELS = [
-    ModelConfig(
-        name="Toxic-BERT",
-        description="Fine-tuned for toxic content detection",
-        type=ModelType.LOCAL,
-        model_path="unitary/toxic-bert"
-    )
+    {
+        "name": "Toxic-BERT",
+        "description": "Fine-tuned for toxic content detection",
+        "type": LOCAL,
+        "model_path": "unitary/toxic-bert"
+    }
 ]
 
-class LocalModelManager:
-    def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {self.device}")
-        self.models = {}
-        self.tokenizers = {}
-        self.pipelines = {}
-        
-    def preload_models(self):
-        """起動時にすべてのローカルモデルを事前にロード"""
-        logger.info("Preloading all local models...")
-        for model in TEXT_GENERATION_MODELS:
-            if model.type == ModelType.LOCAL and model.model_path:
-                self.load_model_sync(model.model_path, "text-generation")
-                
-        for model in CLASSIFICATION_MODELS:
-            if model.type == ModelType.LOCAL and model.model_path:
-                self.load_model_sync(model.model_path, "text-classification")
-        
-        logger.info("All local models preloaded successfully")
+# グローバル変数でモデルやトークナイザーを管理
+models = {}
+tokenizers = {}
+pipelines = {}
+api_clients = {}
+
+# インファレンスAPIクライアントの初期化
+def initialize_api_clients():
+    """Inference APIクライアントの初期化"""
+    for model in TEXT_GENERATION_MODELS + CLASSIFICATION_MODELS:
+        if model["type"] == INFERENCE_API and "model_id" in model:
+            api_clients[model["model_id"]] = AsyncInferenceClient(
+                model["model_id"],
+                token=True  # これによりHFトークンを使用
+            )
+    logger.info("API clients initialized")
+
+# モデルのロード関数
+def load_model(model_path, task="text-generation"):
+    """モデルの同期ロード"""
+    if model_path not in models:
+        logger.info(f"Loading model: {model_path}")
+        try:
+            tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
+            
+            if task == "text-generation":
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_path,
+                    torch_dtype=torch.float16,
+                    load_in_8bit=True,  # メモリ使用量削減のため8bit量子化を使用
+                    device_map="auto"
+                )
+                pipelines[model_path] = pipeline(
+                    "text-generation",
+                    model=model,
+                    tokenizer=tokenizers[model_path]
+                )
+            else:  # classification
+                model = AutoModelForSequenceClassification.from_pretrained(
+                    model_path,
+                    device_map="auto"
+                )
+                pipelines[model_path] = pipeline(
+                    "text-classification",
+                    model=model,
+                    tokenizer=tokenizers[model_path]
+                )
+            
+            models[model_path] = model
+            logger.info(f"Model loaded successfully: {model_path}")
+        except Exception as e:
+            logger.error(f"Error loading model {model_path}: {str(e)}")
+            raise
 
-    def load_model_sync(self, model_path: str, task: str = "text-generation"):
-        """モデルの同期ロード"""
-        if model_path not in self.models:
-            logger.info(f"Loading model: {model_path}")
+# すべてのモデルを事前にロード
+def preload_models():
+    """起動時にすべてのローカルモデルを事前にロード"""
+    logger.info("Preloading all local models...")
+    for model in TEXT_GENERATION_MODELS:
+        if model["type"] == LOCAL and "model_path" in model:
             try:
-                self.tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
-                
-                if task == "text-generation":
-                    model = AutoModelForCausalLM.from_pretrained(
-                        model_path,
-                        torch_dtype=torch.float16,
-                        device_map="auto"
-                    )
-                    self.pipelines[model_path] = pipeline(
-                        "text-generation",
-                        model=model,
-                        tokenizer=self.tokenizers[model_path]
-                    )
-                else:  # classification
-                    model = AutoModelForSequenceClassification.from_pretrained(
-                        model_path,
-                        device_map="auto"
-                    )
-                    self.pipelines[model_path] = pipeline(
-                        "text-classification",
-                        model=model,
-                        tokenizer=self.tokenizers[model_path]
-                    )
-                
-                self.models[model_path] = model
-                logger.info(f"Model loaded successfully: {model_path}")
+                load_model(model["model_path"], "text-generation")
             except Exception as e:
-                logger.error(f"Error loading model {model_path}: {str(e)}")
-                raise
-
-    async def load_model(self, model_path: str, task: str = "text-generation"):
-        """モデルの遅延ロード（バックワードコンパティビリティのために維持）"""
-        if model_path not in self.models:
-            self.load_model_sync(model_path, task)
-
-    @spaces.GPU()
-    def _generate_text_sync(self, pipeline, text: str) -> str:
-        """同期的なテキスト生成の実行"""
-        outputs = pipeline(
+                logger.error(f"Failed to preload {model['name']}: {e}")
+            
+    for model in CLASSIFICATION_MODELS:
+        if model["type"] == LOCAL and "model_path" in model:
+            try:
+                load_model(model["model_path"], "text-classification")
+            except Exception as e:
+                logger.error(f"Failed to preload {model['name']}: {e}")
+    
+    logger.info("Model preloading complete")
+
+# テキスト生成の実行関数
+@spaces.GPU()
+def generate_text(model_path, text):
+    """テキスト生成の実行"""
+    if model_path not in models:
+        load_model(model_path, "text-generation")
+    
+    try:
+        outputs = pipelines[model_path](
             text,
-            max_new_tokens=100,
+            max_new_tokens=50,  # トークン数を減らしてGPUメモリ使用量を削減
             do_sample=True,
             temperature=0.7,
             top_p=0.9,
             num_return_sequences=1
         )
         return outputs[0]["generated_text"]
-
-    async def generate_text(self, model_path: str, text: str) -> str:
-        """テキスト生成の実行（非同期ラッパー）"""
-        if model_path not in self.models:
-            await self.load_model(model_path, "text-generation")
-        
-        try:
-            return self._generate_text_sync(self.pipelines[model_path], text)
-        except Exception as e:
-            logger.error(f"Error in text generation with {model_path}: {str(e)}")
-            raise
-
-    @spaces.GPU()
-    def _classify_text_sync(self, pipeline, text: str) -> str:
-        """同期的なテキスト分類の実行"""
-        result = pipeline(text)
-        return str(result)
-
-    async def classify_text(self, model_path: str, text: str) -> str:
-        """テキスト分類の実行（非同期ラッパー）"""
-        if model_path not in self.models:
-            await self.load_model(model_path, "text-classification")
-        
-        try:
-            return self._classify_text_sync(self.pipelines[model_path], text)
-        except Exception as e:
-            logger.error(f"Error in classification with {model_path}: {str(e)}")
-            raise
-
-class ModelManager:
-    def __init__(self):
-        self.api_clients = {}
-        self.local_manager = LocalModelManager()
-        self._initialize_clients()
-
-    def _initialize_clients(self):
-        """Inference APIクライアントの初期化"""
-        for model in TEXT_GENERATION_MODELS + CLASSIFICATION_MODELS:
-            if model.type == ModelType.INFERENCE_API and model.model_id:
-                self.api_clients[model.model_id] = AsyncInferenceClient(
-                    model.model_id,
-                    token=True  # これによりHFトークンを使用
-                )
+    except Exception as e:
+        logger.error(f"Error in text generation with {model_path}: {str(e)}")
+        raise
+
+# テキスト分類の実行関数
+@spaces.GPU()
+def classify_text(model_path, text):
+    """テキスト分類の実行"""
+    if model_path not in models:
+        load_model(model_path, "text-classification")
     
-    def preload_models(self):
-        """起動時にすべてのモデルを事前にロード"""
-        logger.info("Preloading models...")
-        self.local_manager.preload_models()
-        logger.info("Model preloading complete")
-
-    async def run_text_generation(self, text: str, selected_types: List[str]) -> List[str]:
-        """テキスト生成モデルの実行"""
-        results = []
-        for model in TEXT_GENERATION_MODELS:
-            if model.type.value in selected_types:
-                try:
-                    if model.type == ModelType.INFERENCE_API:
-                        logger.info(f"Running API text generation: {model.name}")
-                        response = await self.api_clients[model.model_id].text_generation(
-                            text, max_new_tokens=100, temperature=0.7
-                        )
-                        results.append(f"{model.name}: {response}")
-                    else:
-                        logger.info(f"Running local text generation: {model.name}")
-                        response = await self.local_manager.generate_text(model.model_path, text)
-                        results.append(f"{model.name}: {response}")
-                except Exception as e:
-                    logger.error(f"Error in {model.name}: {str(e)}")
-                    results.append(f"{model.name}: Error - {str(e)}")
-        return results
-
-    async def run_classification(self, text: str, selected_types: List[str]) -> List[str]:
-        """分類モデルの実行"""
-        results = []
-        for model in CLASSIFICATION_MODELS:
-            if model.type.value in selected_types:
-                try:
-                    if model.type == ModelType.INFERENCE_API:
-                        logger.info(f"Running API classification: {model.name}")
-                        response = await self.api_clients[model.model_id].text_classification(text)
-                        results.append(f"{model.name}: {response}")
-                    else:
-                        logger.info(f"Running local classification: {model.name}")
-                        response = await self.local_manager.classify_text(model.model_path, text)
-                        results.append(f"{model.name}: {response}")
-                except Exception as e:
-                    logger.error(f"Error in {model.name}: {str(e)}")
-                    results.append(f"{model.name}: Error - {str(e)}")
-        return results
-
-class UIComponents:
-    def __init__(self):
-        self.input_text = None
-        self.filter_checkboxes = None
-        self.invoke_button = None
-        self.gen_model_outputs = []
-        self.class_model_outputs = []
-        self.community_output = None
-
-    def create_header(self):
-        """ヘッダーセクションの作成"""
-        return gr.Markdown("""
-        # Toxic Eye
-        This system evaluates the toxicity level of input text using multiple approaches.
-        """)
-
-    def create_input_section(self):
-        """入力セクションの作成"""
-        with gr.Row():
-            self.input_text = gr.Textbox(
-                label="Input Text",
-                placeholder="Enter text to analyze...",
-                lines=3
-            )
-
-    def create_filter_section(self):
-        """フィルターセクションの作成"""
-        with gr.Row():
-            self.filter_checkboxes = gr.CheckboxGroup(
-                choices=[t.value for t in ModelType],
-                value=[t.value for t in ModelType],
-                label="Filter Models",
-                info="Choose which types of models to display",
-                interactive=True
-            )
-
-    def create_invoke_button(self):
-        """Invokeボタンの作成"""
-        with gr.Row():
-            self.invoke_button = gr.Button(
-                "Invoke Selected Models",
-                variant="primary",
-                size="lg"
-            )
-
-    def create_model_grid(self, models: List[ModelConfig]) -> List[Dict]:
-        """モデルグリッドの作成"""
-        outputs = []
-        with gr.Column() as container:
-            for i in range(0, len(models), 2):
-                with gr.Row() as row:
-                    for j in range(min(2, len(models) - i)):
-                        model = models[i + j]
-                        with gr.Column():
-                            with gr.Group() as group:
-                                gr.Markdown(f"### {model.name}")
-                                gr.Markdown(f"Type: {model.type.value}")
-                                output = gr.Textbox(
-                                    label="Model Output",
-                                    lines=5,
-                                    interactive=False,
-                                    info=model.description
-                                )
-                                outputs.append({
-                                    "type": model.type.value,
-                                    "name": model.name,
-                                    "output": output,
-                                    "group": group
-                                })
-        return outputs
-
-    def create_model_tabs(self):
-        """モデルタブの作成"""
-        with gr.Tabs():
-            with gr.Tab("Text Generation LLM"):
-                self.gen_model_outputs = self.create_model_grid(TEXT_GENERATION_MODELS)
-            with gr.Tab("Classification LLM"):
-                self.class_model_outputs = self.create_model_grid(CLASSIFICATION_MODELS)
-            with gr.Tab("Community (Not implemented)"):
-                with gr.Column():
-                    self.community_output = gr.Textbox(
-                        label="Related Community Topics",
-                        lines=5,
-                        interactive=False
+    try:
+        result = pipelines[model_path](text)
+        return str(result)
+    except Exception as e:
+        logger.error(f"Error in classification with {model_path}: {str(e)}")
+        raise
+
+# 複数のモデルでテキスト生成を実行
+async def run_text_generation(text, selected_types):
+    """テキスト生成モデルの実行"""
+    results = []
+    for model in TEXT_GENERATION_MODELS:
+        if model["type"] in selected_types:
+            try:
+                if model["type"] == INFERENCE_API:
+                    logger.info(f"Running API text generation: {model['name']}")
+                    response = await api_clients[model["model_id"]].text_generation(
+                        text, max_new_tokens=50, temperature=0.7
                     )
-
-class ToxicityApp:
-    def __init__(self):
-        self.ui = UIComponents()
-        self.model_manager = ModelManager()
-        self.models_loaded = False
-
-    def update_model_visibility(self, selected_types: List[str]) -> List[gr.update]:
-        """モデルの表示状態を更新"""
-        logger.info(f"Updating visibility for types: {selected_types}")
-        
-        updates = []
-        for outputs in [self.ui.gen_model_outputs, self.ui.class_model_outputs]:
-            for output in outputs:
-                visible = output["type"] in selected_types
-                logger.info(f"Model {output['name']} (type: {output['type']}): visible = {visible}")
-                updates.append(gr.update(visible=visible))
-        return updates
-
-    async def handle_invoke(self, text: str, selected_types: List[str]) -> List[str]:
-        """Invokeボタンのハンドラ"""
-        gen_results = await self.model_manager.run_text_generation(text, selected_types)
-        class_results = await self.model_manager.run_classification(text, selected_types)
-        
-        # 結果リストの長さを調整
-        gen_results.extend([""] * (len(TEXT_GENERATION_MODELS) - len(gen_results)))
-        class_results.extend([""] * (len(CLASSIFICATION_MODELS) - len(class_results)))
-        
-        return gen_results + class_results
+                    results.append(f"{model['name']}: {response}")
+                else:
+                    logger.info(f"Running local text generation: {model['name']}")
+                    response = generate_text(model["model_path"], text)
+                    results.append(f"{model['name']}: {response}")
+            except Exception as e:
+                logger.error(f"Error in {model['name']}: {str(e)}")
+                results.append(f"{model['name']}: Error - {str(e)}")
+    return results
+
+# 複数のモデルでテキスト分類を実行
+async def run_classification(text, selected_types):
+    """分類モデルの実行"""
+    results = []
+    for model in CLASSIFICATION_MODELS:
+        if model["type"] in selected_types:
+            try:
+                if model["type"] == INFERENCE_API:
+                    logger.info(f"Running API classification: {model['name']}")
+                    response = await api_clients[model["model_id"]].text_classification(text)
+                    results.append(f"{model['name']}: {response}")
+                else:
+                    logger.info(f"Running local classification: {model['name']}")
+                    response = classify_text(model["model_path"], text)
+                    results.append(f"{model['name']}: {response}")
+            except Exception as e:
+                logger.error(f"Error in {model['name']}: {str(e)}")
+                results.append(f"{model['name']}: Error - {str(e)}")
+    return results
+
+# Invokeボタンのハンドラ
+async def handle_invoke(text, selected_types):
+    """Invokeボタンのハンドラ"""
+    gen_results = await run_text_generation(text, selected_types)
+    class_results = await run_classification(text, selected_types)
     
-    def load_models_and_update_ui(self):
-        """モデルをロードしUIを更新する"""
-        try:
-            # モデルのロード
-            self.model_manager.preload_models()
-            self.models_loaded = True
-            logger.info("Models loaded successfully")
-            # ロード完了メッセージを返して、UIのロード中表示を非表示にする
-            return gr.update(visible=False), gr.update(visible=True)
-        except Exception as e:
-            logger.error(f"Error loading models: {e}")
-            return gr.update(value=f"Error loading models: {e}"), gr.update(visible=False)
+    # 結果リストの長さを調整
+    gen_results.extend([""] * (len(TEXT_GENERATION_MODELS) - len(gen_results)))
+    class_results.extend([""] * (len(CLASSIFICATION_MODELS) - len(class_results)))
+    
+    return gen_results + class_results
 
-    def create_ui(self):
-        """UIの作成"""
-        with gr.Blocks() as demo:
-            # ロード中コンポーネント
-            with gr.Group(visible=True) as loading_group:
-                gr.Markdown("""
-                # Toxic Eye
-                
-                ### Loading models... This may take a few minutes.
-                
-                The application is initializing and preloading all models.
-                Please wait while the models are being loaded...
-                """)
+# モデルの表示状態を更新
+def update_model_visibility(selected_types):
+    """モデルの表示状態を更新"""
+    logger.info(f"Updating visibility for types: {selected_types}")
+    
+    updates = []
+    for model_outputs in [gen_model_outputs, class_model_outputs]:
+        for output in model_outputs:
+            visible = output["type"] in selected_types
+            logger.info(f"Model {output['name']} (type: {output['type']}): visible = {visible}")
+            updates.append(gr.update(visible=visible))
+    return updates
+
+# モデルをロードしUIを更新する
+def load_models_and_update_ui():
+    """モデルをロードしUIを更新する"""
+    try:
+        # APIクライアント初期化
+        initialize_api_clients()
+        # モデルのロード
+        preload_models()
+        logger.info("Models loaded successfully")
+        # ロード完了メッセージを返して、UIのロード中表示を非表示にする
+        return gr.update(visible=False), gr.update(visible=True)
+    except Exception as e:
+        logger.error(f"Error loading models: {e}")
+        return gr.update(value=f"Error loading models: {e}"), gr.update(visible=False)
+
+# モデルグリッドの作成
+def create_model_grid(models):
+    """モデルグリッドの作成"""
+    outputs = []
+    with gr.Column() as container:
+        for i in range(0, len(models), 2):
+            with gr.Row() as row:
+                for j in range(min(2, len(models) - i)):
+                    model = models[i + j]
+                    with gr.Column():
+                        with gr.Group() as group:
+                            gr.Markdown(f"### {model['name']}")
+                            gr.Markdown(f"Type: {model['type']}")
+                            output = gr.Textbox(
+                                label="Model Output",
+                                lines=5,
+                                interactive=False,
+                                info=model['description']
+                            )
+                            outputs.append({
+                                "type": model["type"],
+                                "name": model["name"],
+                                "output": output,
+                                "group": group
+                            })
+    return outputs
+
+# グローバル変数としてUI部品を保持
+input_text = None
+filter_checkboxes = None
+invoke_button = None
+gen_model_outputs = []
+class_model_outputs = []
+community_output = None
+
+# UIの作成
+def create_ui():
+    """UIの作成"""
+    global input_text, filter_checkboxes, invoke_button, gen_model_outputs, class_model_outputs, community_output
+    
+    with gr.Blocks() as demo:
+        # ロード中コンポーネント
+        with gr.Group(visible=True) as loading_group:
+            gr.Markdown("""
+            # Toxic Eye
             
-            # メインUIコンポーネント（初期状態では非表示）
-            with gr.Group(visible=False) as main_ui_group:
-                self.ui.create_header()
-                self.ui.create_input_section()
-                self.ui.create_filter_section()
-                self.ui.create_invoke_button()
-                self.ui.create_model_tabs()
-
-                # イベントハンドラの設定
-                self.ui.filter_checkboxes.change(
-                    fn=self.update_model_visibility,
-                    inputs=[self.ui.filter_checkboxes],
-                    outputs=[
-                        output["group"]
-                        for outputs in [self.ui.gen_model_outputs, self.ui.class_model_outputs]
-                        for output in outputs
-                    ]
+            ### Loading models... This may take a few minutes.
+            
+            The application is initializing and preloading all models.
+            Please wait while the models are being loaded...
+            """)
+        
+        # メインUIコンポーネント（初期状態では非表示）
+        with gr.Group(visible=False) as main_ui_group:
+            # ヘッダー
+            gr.Markdown("""
+            # Toxic Eye
+            This system evaluates the toxicity level of input text using multiple approaches.
+            """)
+            
+            # 入力セクション
+            with gr.Row():
+                input_text = gr.Textbox(
+                    label="Input Text",
+                    placeholder="Enter text to analyze...",
+                    lines=3
                 )
-
-                self.ui.invoke_button.click(
-                    fn=self.handle_invoke,
-                    inputs=[self.ui.input_text, self.ui.filter_checkboxes],
-                    outputs=[
-                        output["output"]
-                        for outputs in [self.ui.gen_model_outputs, self.ui.class_model_outputs]
-                        for output in outputs
-                    ]
+            
+            # フィルターセクション
+            with gr.Row():
+                filter_checkboxes = gr.CheckboxGroup(
+                    choices=[LOCAL, INFERENCE_API],
+                    value=[LOCAL, INFERENCE_API],
+                    label="Filter Models",
+                    info="Choose which types of models to display",
+                    interactive=True
                 )
             
-            # 起動時にモデルロード処理を実行
-            demo.load(
-                fn=self.load_models_and_update_ui,
-                inputs=None,
-                outputs=[loading_group, main_ui_group]
+            # Invokeボタン
+            with gr.Row():
+                invoke_button = gr.Button(
+                    "Invoke Selected Models",
+                    variant="primary",
+                    size="lg"
+                )
+            
+            # モデルタブ
+            with gr.Tabs():
+                with gr.Tab("Text Generation LLM"):
+                    gen_model_outputs = create_model_grid(TEXT_GENERATION_MODELS)
+                with gr.Tab("Classification LLM"):
+                    class_model_outputs = create_model_grid(CLASSIFICATION_MODELS)
+                with gr.Tab("Community (Not implemented)"):
+                    with gr.Column():
+                        community_output = gr.Textbox(
+                            label="Related Community Topics",
+                            lines=5,
+                            interactive=False
+                        )
+            
+            # イベントハンドラの設定
+            filter_checkboxes.change(
+                fn=update_model_visibility,
+                inputs=[filter_checkboxes],
+                outputs=[
+                    output["group"]
+                    for outputs in [gen_model_outputs, class_model_outputs]
+                    for output in outputs
+                ]
             )
+            
+            invoke_button.click(
+                fn=handle_invoke,
+                inputs=[input_text, filter_checkboxes],
+                outputs=[
+                    output["output"]
+                    for outputs in [gen_model_outputs, class_model_outputs]
+                    for output in outputs
+                ]
+            )
+        
+        # 起動時にモデルロード処理を実行
+        demo.load(
+            fn=load_models_and_update_ui,
+            inputs=None,
+            outputs=[loading_group, main_ui_group]
+        )
+    
+    return demo
 
-        return demo
-
+# メイン関数
 def main():
     logger.info("Starting Toxic Eye application")
-    app = ToxicityApp()
-    demo = app.create_ui()
+    demo = create_ui()
     demo.launch()
 
 if __name__ == "__main__":