Upload app.py with huggingface_hub

app.py

@@ -20,15 +20,18 @@ torch.load = _patched_torch_load
 import sys
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
-import lightning as L
+import random
 import numpy as np
 import gradio as gr
 import whisper
 import tempfile
 import time
+import json
+from pathlib import Path
 from transformers import AutoConfig
 from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
 from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
 
 from src.audiointeraction.dataset.TOKENS import (
     ASSISTANT, AUDIO_BEGIN, ENGLISH, KEEP_SILENCE, ONLINE, PAD,
@@ -40,11 +43,6 @@ from src.audiointeraction.generate.base import (
 )
 from src.audiointeraction.model import GPT, Config
 from src.audiointeraction.tokenizer import Tokenizer
-from src.audiointeraction.utils import get_default_supported_precision
-
-from safetensors.torch import load_file
-from pathlib import Path
-import json
 
 # ── Constants ──────────────────────────────────────────────
 
@@ -64,6 +62,15 @@ EMOTION_EMOJI = {
 }
 
 
+def set_seed(seed: int = 1337) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
 # ── Model loading (module scope, eager) ────────────────────
 
 def _resolve_checkpoint_paths(checkpoint_dir: str):
@@ -76,32 +83,6 @@ def _resolve_checkpoint_paths(checkpoint_dir: str):
     )
 
 
-def _load_model(fabric, model_config_dir, trained_checkpoint):
-    """Load GPT from sharded safetensors using Lightning Fabric."""
-    config = Config.from_file(Path(model_config_dir) / "model_config.yaml")
-    with fabric.init_module(empty_init=True):
-        model = GPT(config)
-    model = fabric.setup(model)
-
-    checkpoint_dir = Path(trained_checkpoint)
-    index_path = checkpoint_dir / "model.safetensors.index.json"
-    if not index_path.is_file():
-        raise FileNotFoundError(f"No model.safetensors.index.json under {checkpoint_dir}")
-
-    with open(index_path) as f:
-        index = json.load(f)
-    shard_files = sorted(set(index["weight_map"].values()))
-
-    state_dict = {}
-    for shard in shard_files:
-        state_dict.update(load_file(str(checkpoint_dir / shard), device="cpu"))
-
-    missing, unexpected = model.load_state_dict(state_dict, strict=True)
-    if missing or unexpected:
-        print(f"[load_model] missing={missing[:3]}… unexpected={unexpected[:3]}…")
-    return model
-
-
 def _load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
     cfg = AutoConfig.from_pretrained(qwen_omni_ckpt)
     audio_cfg = cfg.thinker_config.audio_config
@@ -125,15 +106,28 @@ print(f"Model downloaded to {ckpt_dir}")
 model_config_dir, trained_checkpoint, qwen_omni_ckpt, audio_tower_ckpt = \
     _resolve_checkpoint_paths(ckpt_dir)
 
-set_seed_called = False
-# Use Lightning Fabric for proper precision and device management
+set_seed(1337)
 device = "cuda"
-precision = get_default_supported_precision(training=False)
-fabric = L.Fabric(devices=1, precision=precision)
 
 print("Loading language model...")
-model = _load_model(fabric, model_config_dir, trained_checkpoint)
-model = fabric.setup(model)
+config = Config.from_file(Path(model_config_dir) / "model_config.yaml")
+model = GPT(config)
+model.max_seq_length = config.block_size
+
+checkpoint_dir = Path(trained_checkpoint)
+index_path = checkpoint_dir / "model.safetensors.index.json"
+with open(index_path) as f:
+    index = json.load(f)
+shard_files = sorted(set(index["weight_map"].values()))
+state_dict = {}
+for shard in shard_files:
+    state_dict.update(load_file(str(checkpoint_dir / shard), device="cpu"))
+missing, unexpected = model.load_state_dict(state_dict, strict=True)
+if missing or unexpected:
+    print(f"[load_model] missing={missing[:3]}… unexpected={unexpected[:3]}…")
+del state_dict  # free memory
+
+model = model.to(device).to(torch.bfloat16)
 model.eval()
 
 print("Loading audio encoder...")
@@ -250,11 +244,11 @@ def interact(audio_file, text_instruction=""):
                     if int_token == TEXT_BEGIN:
                         listening = False
                         current_text = [int_token]
-                        print(f"[interact] TEXT_BEGIN at chunk {audio_idx}/{n_chunks}")
+                        print(f"[interact] TEXT_BEGIN at chunk {audio_idx}/{n_chunks}", flush=True)
                     elif int_token == KEEP_SILENCE:
-                        print(f"[interact] KEEP_SILENCE at chunk {audio_idx}/{n_chunks}")
+                        print(f"[interact] KEEP_SILENCE at chunk {audio_idx}/{n_chunks}", flush=True)
                     else:
-                        print(f"[interact] Unexpected token {int_token} at chunk {audio_idx}/{n_chunks}")
+                        print(f"[interact] Unexpected token {int_token} at chunk {audio_idx}/{n_chunks}", flush=True)
                         break
                 else:
                     current_text.append(int_token)
@@ -278,7 +272,7 @@ def interact(audio_file, text_instruction=""):
         # Handle any incomplete turn
         if not listening and len(current_text) > 1:
             text_tokens = current_text[1:]  # skip TEXT_BEGIN
-            if text_tokens[0] in EMOTION_EMOJI if text_tokens else False:
+            if text_tokens and text_tokens[0] in EMOTION_EMOJI:
                 emotion_tag = EMOTION_EMOJI[text_tokens[0]]
                 text_tokens = text_tokens[1:]
             decoded = tokenizer.decode(torch.tensor(text_tokens))
@@ -312,7 +306,6 @@ def interact(audio_file, text_instruction=""):
 
 CSS = """
 #col-container { max-width: 900px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
 """
 
 with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo: