Add 8b model

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 import pandas as pd
 import numpy as np
 import torch
-import subprocess
 from threading import Thread
 from transformers import (
     AutoModelForCausalLM,
@@ -13,21 +12,62 @@ from transformers import (
 )
 
 # ─── MODEL SETUP ────────────────────────────────────────────────────────────────
-MODEL_NAME = "bytedance-research/ChatTS-14B"
+# Default to 8B but keep both variants resident on the GPU.
+DEFAULT_MODEL_NAME = "bytedance-research/ChatTS-8B"
+AVAILABLE_MODEL_NAMES = [
+    "bytedance-research/ChatTS-8B",
+    "bytedance-research/ChatTS-14B"
+]
+
+MODEL_REGISTRY = {}
+
+for name in AVAILABLE_MODEL_NAMES:
+    print(f"Loading model into memory: {name}")
+    tok = AutoTokenizer.from_pretrained(name, trust_remote_code=True)
+    proc = AutoProcessor.from_pretrained(name, trust_remote_code=True, tokenizer=tok)
+    mdl = AutoModelForCausalLM.from_pretrained(
+        name,
+        trust_remote_code=True,
+        device_map="auto",
+        torch_dtype=torch.float16
+    )
+    mdl.eval()
+    MODEL_REGISTRY[name] = {
+        "tokenizer": tok,
+        "processor": proc,
+        "model": mdl
+    }
 
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_NAME, trust_remote_code=True
-)
-processor = AutoProcessor.from_pretrained(
-    MODEL_NAME, trust_remote_code=True, tokenizer=tokenizer
-)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    trust_remote_code=True,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
-model.eval()
+CURRENT_MODEL_NAME = DEFAULT_MODEL_NAME
+
+tokenizer = MODEL_REGISTRY[CURRENT_MODEL_NAME]["tokenizer"]
+processor = MODEL_REGISTRY[CURRENT_MODEL_NAME]["processor"]
+model = MODEL_REGISTRY[CURRENT_MODEL_NAME]["model"]
+
+
+def load_model_by_name(name: str):
+    """Activate the preloaded model by name without reloading weights."""
+    global tokenizer, processor, model, CURRENT_MODEL_NAME
+
+    if name not in MODEL_REGISTRY:
+        return f"Model not available: {name}"
+
+    if name == CURRENT_MODEL_NAME:
+        return f"Model already selected: {name}"
+
+    CURRENT_MODEL_NAME = name
+    tokenizer = MODEL_REGISTRY[name]["tokenizer"]
+    processor = MODEL_REGISTRY[name]["processor"]
+    model = MODEL_REGISTRY[name]["model"]
+    model.eval()
+
+    print(f"Activated model: {name}")
+    return f"Active model: {name}"
+
+
+def switch_model(selected_model_name: str):
+    """Wrapper for Gradio to switch models; returns status text."""
+    return load_model_by_name(selected_model_name)
 
 # ─── HELPER FUNCTIONS ──────────────────────────────────────────────────────────
 
@@ -290,6 +330,21 @@ with gr.Blocks(title="ChatTS Demo") as demo:
 
     with gr.Row():
         with gr.Column(scale=1):
+            # Model selection UI
+            model_radio = gr.Radio(
+                choices=["bytedance-research/ChatTS-8B", "bytedance-research/ChatTS-14B"],
+                value=CURRENT_MODEL_NAME,
+                label="Model Version"
+            )
+
+            model_btn = gr.Button("Load Model")
+
+            model_status = gr.Textbox(
+                label="Model Status",
+                value=f"Models in memory: {', '.join(AVAILABLE_MODEL_NAMES)}. Active: {CURRENT_MODEL_NAME}",
+                interactive=False
+            )
+
             upload = gr.File(
                 label="Upload CSV File",
                 file_types=[".csv"],
@@ -355,5 +410,12 @@ with gr.Blocks(title="ChatTS Demo") as demo:
         outputs=[text_out]
     )
 
+    # Wire model loading button
+    model_btn.click(
+        fn=switch_model,
+        inputs=[model_radio],
+        outputs=[model_status]
+    )
+
 if __name__ == '__main__':
     demo.launch()

app_legacy.py

@@ -1,53 +0,0 @@
-import spaces                        # for ZeroGPU support
-import gradio as gr
-import pandas as pd
-import numpy as np
-import torch
-import subprocess
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    AutoProcessor,
-)
-
-# ─── MODEL SETUP ────────────────────────────────────────────────────────────────
-MODEL_NAME = "bytedance-research/ChatTS-14B"
-
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_NAME, trust_remote_code=True
-)
-processor = AutoProcessor.from_pretrained(
-    MODEL_NAME, trust_remote_code=True, tokenizer=tokenizer
-)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    trust_remote_code=True,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
-model.eval()
-
-
-# ─── INFERENCE + VALIDATION ────────────────────────────────────────────────────
-@spaces.GPU
-def generate_text(prompt):
-    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs, 
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.2,
-        top_p=0.9
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-demo = gr.Interface(
-    fn=generate_text,
-    inputs=gr.Textbox(lines=2, label="Prompt"),
-    outputs=gr.Textbox(lines=6, label="Generated Text")
-)
-
-
-if __name__ == '__main__':
-    subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
-    demo.launch()