Update app.py

app.py

@@ -1,166 +1,166 @@
-import gradio as gr
-import torch
-import os
-import sys
-import warnings
-from pathlib import Path
-
-# Add root to path to allow imports from project root when running from demo-code/
-# or when running from root
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.dirname(current_dir)
-sys.path.append(current_dir)
-sys.path.append(parent_dir)
-
-# Import project modules
-try:
-    from inference import load_trained_model, inference as run_inference_cmd
-    from visualize import visualize
-    from model import setup_model_and_tokenizer, get_motion_token_info
-    from generate import generate_t2m
-    from data import compute_length_stats, build_prompt_vocab, check_has_participant_id, load_dataset
-    import config
-except ImportError as e:
-    print(f"Error importing project modules: {e}")
-    print("Make sure you are running this from the project root or have the project structure intact.")
-
-# Constants
-HF_REPO_ID = "rdz-falcon/SignMotionGPT"
-EPOCH_SUBFOLDER = "stage2/epoch-030" 
-
-def load_model_from_hf(repo_id, subfolder, token=None):
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    print(f"Loading model from HF: {repo_id}/{subfolder}")
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder=subfolder, token=token, trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained(repo_id, subfolder=subfolder, token=token, trust_remote_code=True)
-        return model, tokenizer
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None, None
-
-# Global model cache
-MODEL = None
-TOKENIZER = None
-MOTION_TOKEN_IDS = None
-MOT_BEGIN_ID = None
-MOT_END_ID = None
-CODEBOOK_SIZE = 512 
-
-def init_model():
-    global MODEL, TOKENIZER, MOTION_TOKEN_IDS, MOT_BEGIN_ID, MOT_END_ID
-    if MODEL is not None:
-        return
-
-    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
-    
-    # Load model/tokenizer
-    MODEL, TOKENIZER = load_model_from_hf(HF_REPO_ID, EPOCH_SUBFOLDER, token)
-    
-    if MODEL is None:
-        raise RuntimeError(f"Failed to load model from {HF_REPO_ID}/{EPOCH_SUBFOLDER}")
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    MODEL.to(device)
-    MODEL.eval()
-
-    # Setup token info
-    motion_token_ids = []
-    for i in range(CODEBOOK_SIZE):
-        t = f"<motion_{i}>"
-        if t in TOKENIZER.get_vocab():
-            motion_token_ids.append(TOKENIZER.convert_tokens_to_ids(t))
-    
-    MOTION_TOKEN_IDS = motion_token_ids
-    MOT_BEGIN_ID = TOKENIZER.convert_tokens_to_ids("<MOT_BEGIN>") if "<MOT_BEGIN>" in TOKENIZER.get_vocab() else None
-    MOT_END_ID = TOKENIZER.convert_tokens_to_ids("<MOT_END>") if "<MOT_END>" in TOKENIZER.get_vocab() else None
-    
-    print("Model initialized.")
-
-def generate_motion_app(text_prompt):
-    if not text_prompt:
-        return None, "Please enter a prompt."
-    
-    if MODEL is None:
-        try:
-            init_model()
-        except Exception as e:
-            return None, f"Model Initialization Failed: {e}"
-        
-    device = MODEL.device
-    print(f"Generating for: {text_prompt}")
-    
-    try:
-        generated_tokens = generate_t2m(
-            model=MODEL,
-            tokenizer=TOKENIZER,
-            prompt_text=text_prompt,
-            mot_begin_id=MOT_BEGIN_ID,
-            mot_end_id=MOT_END_ID,
-            motion_token_ids=MOTION_TOKEN_IDS,
-            length_stats_by_text={}, # Fallback to global_median_len
-            global_median_len=100,   # Reasonable default
-            prompt_vocab=None,
-            has_pid=False,
-            per_prompt_vocab=False   # Allow all tokens
-        )
-    except Exception as e:
-        return None, f"Generation Error: {e}"
-
-    # Visualization
-    try:
-        # Ensure paths for VQ-VAE and SMPL-X
-        # In HF Spaces, we assume these are in the repo (e.g., ./data)
-        data_dir = os.environ.get("DATA_DIR", "data")
-        vqvae_ckpt = os.path.join(data_dir, "vqvae_model.pt")
-        stats_path = os.path.join(data_dir, "vqvae_stats.pt")
-        smplx_dir = os.path.join(data_dir, "smplx_models")
-        
-        # Check existence
-        missing = []
-        if not os.path.exists(vqvae_ckpt): missing.append(vqvae_ckpt)
-        if not os.path.exists(stats_path): missing.append(stats_path)
-        if not os.path.exists(smplx_dir): missing.append(smplx_dir)
-        
-        if missing:
-            return None, f"Missing visualization files in {data_dir}: {missing}. Please ensure they are uploaded to the Space."
-
-        # Output to a temporary file
-        # Gradio needs a file path or HTML string. visualize returns a Figure.
-        output_html = "temp_viz.html"
-        
-        fig = visualize(
-            tokens=generated_tokens,
-            vqvae_ckpt=vqvae_ckpt,
-            stats_path=stats_path,
-            smplx_dir=smplx_dir,
-            output_html=output_html,
-            title=f"Motion: {text_prompt}",
-            fps=20
-        )
-        
-        if fig is None:
-             return None, "Visualization failed (no frames produced)."
-             
-        return fig, f"Success! Generated tokens length: {len(generated_tokens.split())}"
-        
-    except Exception as e:
-        return None, f"Visualization Error: {e}"
-
-
-# Gradio UI
-with gr.Interface(
-    fn=generate_motion_app,
-    inputs=gr.Textbox(label="Enter Motion Prompt", placeholder="e.g. walking forward"),
-    outputs=[
-        gr.Plot(label="Motion Visualization"),
-        gr.Textbox(label="Status/Output")
-    ],
-    title="SignMotionGPT Demo",
-    description="Generate Sign Language/Motion Avatars from Text. Using model checkpoint: epoch 30."
-) as demo:
-    pass
-
-if __name__ == "__main__":
-    demo.launch()
-
+import gradio as gr
+import torch
+import os
+import sys
+import warnings
+from pathlib import Path
+
+# Add root to path to allow imports from project root when running from demo-code/
+# or when running from root
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.append(current_dir)
+sys.path.append(parent_dir)
+
+# Import project modules
+try:
+    from inference import load_trained_model, inference as run_inference_cmd
+    from visualize import visualize
+    from model import setup_model_and_tokenizer, get_motion_token_info
+    from generate import generate_t2m
+    from data import compute_length_stats, build_prompt_vocab, check_has_participant_id, load_dataset
+    import config
+except ImportError as e:
+    print(f"Error importing project modules: {e}")
+    print("Make sure you are running this from the project root or have the project structure intact.")
+
+# Constants
+HF_REPO_ID = "rdz-falcon/SignMotionGPTfit-archive"
+EPOCH_SUBFOLDER = "stage2/epoch-030" 
+
+def load_model_from_hf(repo_id, subfolder, token=None):
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    print(f"Loading model from HF: {repo_id}/{subfolder}")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder=subfolder, token=token, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(repo_id, subfolder=subfolder, token=token, trust_remote_code=True)
+        return model, tokenizer
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None, None
+
+# Global model cache
+MODEL = None
+TOKENIZER = None
+MOTION_TOKEN_IDS = None
+MOT_BEGIN_ID = None
+MOT_END_ID = None
+CODEBOOK_SIZE = 512 
+
+def init_model():
+    global MODEL, TOKENIZER, MOTION_TOKEN_IDS, MOT_BEGIN_ID, MOT_END_ID
+    if MODEL is not None:
+        return
+
+    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+    
+    # Load model/tokenizer
+    MODEL, TOKENIZER = load_model_from_hf(HF_REPO_ID, EPOCH_SUBFOLDER, token)
+    
+    if MODEL is None:
+        raise RuntimeError(f"Failed to load model from {HF_REPO_ID}/{EPOCH_SUBFOLDER}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    MODEL.to(device)
+    MODEL.eval()
+
+    # Setup token info
+    motion_token_ids = []
+    for i in range(CODEBOOK_SIZE):
+        t = f"<motion_{i}>"
+        if t in TOKENIZER.get_vocab():
+            motion_token_ids.append(TOKENIZER.convert_tokens_to_ids(t))
+    
+    MOTION_TOKEN_IDS = motion_token_ids
+    MOT_BEGIN_ID = TOKENIZER.convert_tokens_to_ids("<MOT_BEGIN>") if "<MOT_BEGIN>" in TOKENIZER.get_vocab() else None
+    MOT_END_ID = TOKENIZER.convert_tokens_to_ids("<MOT_END>") if "<MOT_END>" in TOKENIZER.get_vocab() else None
+    
+    print("Model initialized.")
+
+def generate_motion_app(text_prompt):
+    if not text_prompt:
+        return None, "Please enter a prompt."
+    
+    if MODEL is None:
+        try:
+            init_model()
+        except Exception as e:
+            return None, f"Model Initialization Failed: {e}"
+        
+    device = MODEL.device
+    print(f"Generating for: {text_prompt}")
+    
+    try:
+        generated_tokens = generate_t2m(
+            model=MODEL,
+            tokenizer=TOKENIZER,
+            prompt_text=text_prompt,
+            mot_begin_id=MOT_BEGIN_ID,
+            mot_end_id=MOT_END_ID,
+            motion_token_ids=MOTION_TOKEN_IDS,
+            length_stats_by_text={}, # Fallback to global_median_len
+            global_median_len=100,   # Reasonable default
+            prompt_vocab=None,
+            has_pid=False,
+            per_prompt_vocab=False   # Allow all tokens
+        )
+    except Exception as e:
+        return None, f"Generation Error: {e}"
+
+    # Visualization
+    try:
+        # Ensure paths for VQ-VAE and SMPL-X
+        # In HF Spaces, we assume these are in the repo (e.g., ./data)
+        data_dir = os.environ.get("DATA_DIR", "data")
+        vqvae_ckpt = os.path.join(data_dir, "vqvae_model.pt")
+        stats_path = os.path.join(data_dir, "vqvae_stats.pt")
+        smplx_dir = os.path.join(data_dir, "smplx_models")
+        
+        # Check existence
+        missing = []
+        if not os.path.exists(vqvae_ckpt): missing.append(vqvae_ckpt)
+        if not os.path.exists(stats_path): missing.append(stats_path)
+        if not os.path.exists(smplx_dir): missing.append(smplx_dir)
+        
+        if missing:
+            return None, f"Missing visualization files in {data_dir}: {missing}. Please ensure they are uploaded to the Space."
+
+        # Output to a temporary file
+        # Gradio needs a file path or HTML string. visualize returns a Figure.
+        output_html = "temp_viz.html"
+        
+        fig = visualize(
+            tokens=generated_tokens,
+            vqvae_ckpt=vqvae_ckpt,
+            stats_path=stats_path,
+            smplx_dir=smplx_dir,
+            output_html=output_html,
+            title=f"Motion: {text_prompt}",
+            fps=20
+        )
+        
+        if fig is None:
+             return None, "Visualization failed (no frames produced)."
+             
+        return fig, f"Success! Generated tokens length: {len(generated_tokens.split())}"
+        
+    except Exception as e:
+        return None, f"Visualization Error: {e}"
+
+
+# Gradio UI
+with gr.Interface(
+    fn=generate_motion_app,
+    inputs=gr.Textbox(label="Enter Motion Prompt", placeholder="e.g. walking forward"),
+    outputs=[
+        gr.Plot(label="Motion Visualization"),
+        gr.Textbox(label="Status/Output")
+    ],
+    title="SignMotionGPT Demo",
+    description="Generate Sign Language/Motion Avatars from Text. Using model checkpoint: epoch 30."
+) as demo:
+    pass
+
+if __name__ == "__main__":
+    demo.launch()
+