Fix render

README.md

@@ -8,21 +8,32 @@ sdk_version: "6.1.0"
 app_file: app_new.py
 pinned: false
 python_version: "3.10"
-short_description: Text-to-3D motion generation using ONNX INT8 models
+short_description: Text-to-3D motion generation using ONNX models
 ---
 
 # MoMask: Text-to-Motion Generation
 
 Generate 3D human skeleton animations from text descriptions using [MoMask](https://github.com/EricGuo5513/momask-codes).
 
-## Model Architecture (ONNX INT8, ~151MB total)
-| Model | Size | Precision |
-|-------|------|-----------|
-| CLIP Text Encoder | 62MB | INT8 |
-| Mask Transformer | 20MB | INT8 |
-| Residual Transformer | 20MB | INT8 |
-| VQ-VAE Decoder | 44MB | FP32 |
-| Length Estimator | 441KB | INT8 |
+## Features
+- Text-to-motion generation with classifier-free guidance
+- Download BVH files for Blender import
+- ~7 seconds of motion per generation
+
+## Model Architecture (ONNX FP32, ~416MB total)
+| Model | Size | Purpose |
+|-------|------|---------|
+| CLIP Text Encoder | 254MB | Text embedding |
+| Mask Transformer | 56MB | Initial motion tokens |
+| Residual Transformer | 55MB | Refine motion details |
+| VQ-VAE Decoder | 46MB | Decode to motion |
+| Length Estimator | 0.5MB | Predict motion length |
 
 ## Usage
-Enter a text description and click Generate.
+1. Enter a text description (e.g., "A person walks forward")
+2. Optionally set duration and seed
+3. Click Generate
+4. Download MP4 video or BVH for Blender
+
+## Credits
+Based on [MoMask](https://github.com/EricGuo5513/momask-codes) by Chuan Guo et al.

app_new.py

@@ -44,7 +44,8 @@ ONNX_DIR = Path(__file__).parent / "onnx_models"
 DEVICE = "cpu"
 JOINTS_NUM = 22
 TIMESTEPS = 18
-COND_SCALE = 4
+MASK_COND_SCALE = 4.0  # CFG scale for mask transformer
+RES_COND_SCALE = 5.0   # CFG scale for residual transformer
 TEMPERATURE = 1.0
 TOPK_FILTER = 0.9
 
@@ -226,8 +227,8 @@ def gumbel_sample(logits, temperature=1.0):
     gumbels = -torch.log(-torch.log(torch.rand_like(logits) + 1e-8) + 1e-8)
     return ((logits / max(temperature, 1e-10)) + gumbels).argmax(dim=-1)
 # ============ Main Generation Pipeline ============
-def generate_motion(text, motion_length=0, seed=None, export_bvh=False):
-    """Generate motion from text prompt"""
+def generate_motion(text, motion_length=0, seed=None):
+    """Generate motion from text prompt with CFG"""
     if seed is not None:
         torch.manual_seed(seed)
         np.random.seed(seed)
@@ -239,6 +240,7 @@ def generate_motion(text, motion_length=0, seed=None, export_bvh=False):
 
     clip_sess = get_session("clip_text")
     text_emb = clip_sess.run(None, {"text_tokens": tokens.numpy()})[0]
+    zero_emb = np.zeros_like(text_emb)  # For CFG unconditional path
 
     if motion_length <= 0:
         len_sess = get_session("length_estimator")
@@ -279,11 +281,18 @@ def generate_motion(text, motion_length=0, seed=None, export_bvh=False):
 
         ids[:, :token_len] = torch.where(is_mask, mask_id, ids[:, :token_len])
 
-        logits = mask_sess.run(None, {
+        # CFG: conditional and unconditional logits
+        cond_logits = mask_sess.run(None, {
             "motion_ids": ids.numpy(),
             "cond_vector": text_emb,
             "padding_mask": padding_mask
         })[0]
+        uncond_logits = mask_sess.run(None, {
+            "motion_ids": ids.numpy(),
+            "cond_vector": zero_emb,
+            "padding_mask": padding_mask
+        })[0]
+        logits = uncond_logits + (cond_logits - uncond_logits) * MASK_COND_SCALE
 
         logits = torch.from_numpy(logits)
         logits = logits[:, :512, :token_len]
@@ -316,12 +325,20 @@ def generate_motion(text, motion_length=0, seed=None, export_bvh=False):
 
         q_id = np.array([q], dtype=np.int64)
 
-        logits = res_sess.run(None, {
+        # CFG for residual transformer
+        cond_logits = res_sess.run(None, {
             "motion_codes": history_sum.astype(np.float32),
             "q_id": q_id,
             "cond_vector": text_emb,
             "padding_mask": padding_mask
         })[0]
+        uncond_logits = res_sess.run(None, {
+            "motion_codes": history_sum.astype(np.float32),
+            "q_id": q_id,
+            "cond_vector": zero_emb,
+            "padding_mask": padding_mask
+        })[0]
+        logits = uncond_logits + (cond_logits - uncond_logits) * RES_COND_SCALE
 
         logits = torch.from_numpy(logits)[:, :512, :token_len].permute(0, 2, 1)
         new_ids_q = gumbel_sample(logits, 1.0)
@@ -344,23 +361,20 @@ def generate_motion(text, motion_length=0, seed=None, export_bvh=False):
     video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     plot_3d_motion(video_path, joints, text, fps=20)
 
-    bvh_path = None
-    if export_bvh:
-        bvh_path = tempfile.NamedTemporaryFile(suffix=".bvh", delete=False).name
-        joints_to_bvh(joints, bvh_path, fps=20)
-        print(f"BVH exported: {bvh_path}")
+    bvh_path = tempfile.NamedTemporaryFile(suffix=".bvh", delete=False).name
+    joints_to_bvh(joints, bvh_path, fps=20)
 
     return joints, video_path, bvh_path
 # ============ Gradio Interface ============
 def create_demo():
     import gradio as gr
 
-    def generate_fn(text, length, seed, export_bvh):
+    def generate_fn(text, length, seed):
         if not text or text.strip() == "":
             return None, None
         seed = int(seed) if seed else None
         length = float(length) if length else 0
-        joints, video_path, bvh_path = generate_motion(text, length, seed, export_bvh)
+        joints, video_path, bvh_path = generate_motion(text, length, seed)
         return video_path, bvh_path
 
     with gr.Blocks(title="MoMask") as demo:
@@ -377,28 +391,27 @@ def create_demo():
                                        info="0 = auto-estimate")
                     seed = gr.Number(label="Seed", value=42,
                                     info="For reproducibility")
-                export_bvh = gr.Checkbox(label="Export BVH for Blender", value=True)
                 btn = gr.Button("Generate", variant="primary")
 
             with gr.Column():
                 video = gr.Video(label="Generated Motion")
-                bvh_file = gr.File(label="BVH Download")
+                bvh_file = gr.File(label="BVH Download (for Blender)")
 
         gr.Examples(
             examples=[
-                ["A person walks forward", 0, 42, True],
-                ["A person is running on a treadmill", 0, 123, True],
-                ["A person jumps up and then lands", 0, 456, True],
-                ["A person does a salsa dance", 0, 789, True],
-                ["A person kicks with their right leg", 0, 101, True],
+                ["A person walks forward", 0, 42],
+                ["A person is running on a treadmill", 0, 123],
+                ["A person jumps up and then lands", 0, 456],
+                ["A person does a salsa dance", 0, 789],
+                ["A person kicks with their right leg", 0, 101],
             ],
-            inputs=[text, length, seed, export_bvh],
+            inputs=[text, length, seed],
             outputs=[video, bvh_file],
             fn=generate_fn,
             cache_examples=False,
         )
 
-        btn.click(fn=generate_fn, inputs=[text, length, seed, export_bvh], outputs=[video, bvh_file])
+        btn.click(fn=generate_fn, inputs=[text, length, seed], outputs=[video, bvh_file])
 
     return demo
 
@@ -409,7 +422,7 @@ if __name__ == "__main__":
         length = float(sys.argv[2]) if len(sys.argv) > 2 else 0
         seed = int(sys.argv[3]) if len(sys.argv) > 3 else 42
 
-        joints, video_path, bvh_path = generate_motion(text, length, seed, export_bvh=True)
+        joints, video_path, bvh_path = generate_motion(text, length, seed)
         print(f"Video: {video_path}")
         print(f"BVH: {bvh_path}")
         print(f"Joints shape: {joints.shape}")

onnx_models/clip_text.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:75213356a2cca6a6a69cb4ea45142dff121bf0baf5bdad621e0e68fe68355b6a
-size 64683509
+oid sha256:fee129a5e73595244105a917c8cd6884bd97f04d6a1d09d00b4e715d590fe90e
+size 254389519

onnx_models/mask_transformer.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8154ce0e72fce61b889c86ebde4e08e1976632f943bde9a49117af7d4b9bd95d
-size 20297925
+oid sha256:eb8513f25349c03a7ead2447a2d40d906011ff813905a921a2424544a6e632e9
+size 56169224

onnx_models/residual_transformer.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1f703b17e8ea1c7557d8f3ba3951c0ca7e56c22e3d3514df5d2f35c08bf0a00c
-size 20039878
+oid sha256:931b5b0bf2b1e507233b48d3108fc0ce89bc49d2fb058d8f6b66c6867b554375
+size 55127345

requirements.txt

@@ -1,4 +1,5 @@
-# Minimal requirements for HuggingFace Space (CPU)
+# MoMask HuggingFace Space requirements (CPU)
+gradio>=6.1.0
 onnxruntime>=1.16.0
 torch>=2.0.0
 numpy