.

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-from transformers import pipeline, AutoModelForCausalLM
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import torch
 import logging
 
@@ -11,44 +11,55 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Stores for models and tokenizers
+tokenizers = {}
+pipelines = {}
+
 # Predefined list of models to compare (can be expanded)
 model_options = {
-    "Foundation-Sec-8B": pipeline("text-generation", model="fdtn-ai/Foundation-Sec-8B", torch_dtype=torch.bfloat16),
+    "Foundation-Sec-8B": "fdtn-ai/Foundation-Sec-8B",
 }
 
-@spaces.GPU
-def generate_text_local(model_pipeline, prompt):
-    """Local text generation"""
+# Initialize models at startup
+for model_name, model_path in model_options.items():
     try:
-        # モデル名取得（なければ 'unknown'）
-        model_name = getattr(getattr(model_pipeline, "model", None), "name_or_path", "unknown")
-        logger.info(f"Running local text generation with {model_name}")
-
-        # Move model to GPU (entire pipeline)
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        if hasattr(model_pipeline, "model"):
-            model_pipeline.model = model_pipeline.model.to(device)
-
-        # Record device information
-        device_info = next(model_pipeline.model.parameters()).device if hasattr(model_pipeline, "model") else "unknown"
-        logger.info(f"Model {model_name} is running on device: {device_info}")
-
-        outputs = model_pipeline(
-            prompt,
-            max_new_tokens=3,
-            do_sample=True,
-            temperature=0.1,
-            top_p=0.9,
-            clean_up_tokenization_spaces=True,
+        logger.info(f"Initializing text generation model: {model_path}")
+        tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
+        pipelines[model_path] = pipeline(
+            "text-generation",
+            model=model_path,
+            tokenizer=tokenizers[model_path],
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True
         )
+        logger.info(f"Model initialized successfully: {model_path}")
+    except Exception as e:
+        logger.error(f"Error initializing model {model_path}: {str(e)}")
 
-        # Move model back to CPU
-        if hasattr(model_pipeline, "model"):
-            model_pipeline.model = model_pipeline.model.to("cpu")
+@spaces.GPU
+def generate_text_local(model_path, prompt, max_new_tokens=512, temperature=0.7, top_p=0.95):
+    """Local text generation"""
+    try:
+        # Use the already initialized model
+        if model_path in pipelines:
+            model_pipeline = pipelines[model_path]
+            logger.info(f"Running text generation with {model_path}")
+
+            outputs = model_pipeline(
+                prompt,
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+                clean_up_tokenization_spaces=True,
+            )
 
-        return outputs[0]["generated_text"].replace(prompt, "").strip()
+            return outputs[0]["generated_text"].replace(prompt, "").strip()
+        else:
+            return f"Error: Model {model_path} not initialized"
     except Exception as e:
-        logger.error(f"Error in local text generation with {model_name}: {str(e)}")
+        logger.error(f"Error in text generation with {model_path}: {str(e)}")
         return f"Error: {str(e)}"
 
 # Build Gradio app
@@ -94,13 +105,23 @@ def create_demo():
         ):
             #if len(selected_models) != 2:
             #    return "Error: Please select exactly two models to compare.", ""
-            responses = generate_text_local(
-                #message, [], system_message, max_tokens, temperature, top_p, selected_models
-                model_options[selected_models[0]],
-                message
+            
+            if len(selected_models) == 0:
+                return "Error: Please select at least one model"
+            
+            model_path = model_options[selected_models[0]]
+            full_prompt = f"{system_message}\n\nUser: {message}\nAssistant:"
+            response = generate_text_local(
+                model_path,
+                full_prompt,
+                max_tokens,
+                temperature,
+                top_p
             )
+            
             #return responses.get(selected_models[0], ""), responses.get(selected_models[1], "")
-            return responses
+            return response
+            
         # Add a button for generating responses
         submit_button = gr.Button("Generate Responses")
         submit_button.click(