[update] revise

.gitignore

@@ -28,3 +28,8 @@ local
 tmp
 .git_bk
 demo
+.git_hf_space
+.git_github/
+*.pyc
+tmp
+*.png

README.md

@@ -20,6 +20,7 @@ short_description: "Generative Refinement Networks"
 ---
 
 ## 🔥 Updates!!
+* June 8, 2026: ✈️ The training & fine-tuning code for GRN-T2I and GRN-T2V is released.
 * June 3, 2026: 🍉 A toy image-video dataset is provided for GRN-T2I/GRN-T2V training and fine-tuning.
 * May 23, 2026: 🌺 We release the training and evaluation code for HBQ tokenizer, enjoy~
 * April 14, 2026: 🤗 Paper and code release
@@ -31,22 +32,22 @@ short_description: "Generative Refinement Networks"
 - [🚀 Demo](#-demo)
 - [📦 Model Zoo](#-model-zoo)
 - [🛠️ Installation](#️-installation)
-- [🖼️ Class-to-Image](#️-class-to-image)
-  - [Dataset](#dataset)
+- [📦 HBQ Tokenizer](#-hbq-tokenizer)
+  - [Data](#data)
   - [Training](#training)
   - [Evaluation](#evaluation)
+- [🖼️ Class-to-Image](#️-class-to-image)
+  - [Data](#data-1)
+  - [Training](#training-2)
+  - [Evaluation](#evaluation-1)
 - [🎨 Text-to-Image](#-text-to-image)
-  - [Data](#data)
-  - [Train](#train)
+  - [Data](#data-2)
+  - [Training](#training-2)
   - [Inference](#inference)
 - [🎬 Text-to-Video](#-text-to-video)
-  - [Data](#data-1)
-  - [Train](#train-1)
+  - [Data](#data-3)
+  - [Training](#training-3)
   - [Inference](#inference-1)
-- [📦 HBQ Tokenizer](#-hbq-tokenizer)
-  - [Data](#data-2)
-  - [Training](#training-1)
-  - [Evaluation](#evaluation-1)
 - [📧 Contact](#-contact)
 - [🤗 Acknowledgements](#-acknowledgements)
 - [📝 Citation](#-citation)
@@ -70,7 +71,7 @@ GRN achieves state-of-the-art results on ImageNet reconstruction and class-condi
 
 <figure align="center">
   <figcaption><strong><em>Generative Refinement Framework</em></strong></figcaption>
-  <img src="demo/framework.jpg" width="100%" alt="Framework">
+  <img src="assets/framework.jpg" width="100%" alt="Framework">
 </figure>
 
 <p align="center">
@@ -130,13 +131,13 @@ Starting from a random token map, GRN randomly selects more predictions at each
 ### GRN-2B Class-to-Image Examples
 <figure align="center">
   <!-- <figcaption><strong><em>GRN-2B Class-to-Image Examples</em></strong></figcaption> -->
-  <img src="demo/c2i_examples.jpg" width="100%" alt="Class-to-Image Examples">
+  <img src="assets/c2i_examples.jpg" width="100%" alt="Class-to-Image Examples">
 </figure>
 
 ### GRN-2B Text-to-Image Examples
 <figure align="center">
   <!-- <figcaption><strong><em>GRN-2B Text-to-Image Examples</em></strong></figcaption> -->
-  <img src="demo/t2i_examples.jpg" width="100%" alt="Text-to-Image Examples">
+  <img src="assets/t2i_examples.jpg" width="100%" alt="Text-to-Image Examples">
 </figure>
 
 ---
@@ -160,7 +161,7 @@ Try our interactive Text-to-Video demo on Discord:
 
 <figure align="center">
   <figcaption><strong><em>T2V Demo on Discord</em></strong></figcaption>
-  <img src="demo/t2v_demo.png" width="100%" alt="T2V Demo">
+  <img src="assets/t2v_demo.png" width="100%" alt="T2V Demo">
 </figure>
 
 ---
@@ -187,8 +188,9 @@ cd GRN
 ### Step 2: Create conda environment
 A suitable [conda](https://conda.io/) environment named `GRN` can be created and activated with:
 ```bash
-conda env create -f environment.yaml
+conda create -n GRN python=3.11
 conda activate GRN
+pip install -r requirements.txt
 ```
 
 ### Troubleshooting
@@ -201,9 +203,44 @@ Check this [issue](https://github.com/conda/conda/issues/13812#issuecomment-2071
 
 ---
 
+## 📦 HBQ Tokenizer
+
+### Data
+Image Dataset, e.g., data_root/username/labels/imagenet/train.txt:
+```
+[image_1_full_path]
+[image_2_full_path]
+[image_3_full_path]
+...
+```
+
+Video Dataset, e.g., data_root/username/labels_hanjian/high-quality-video/horizontal_videos.txt
+```
+[video_1_full_path]
+[video_2_full_path]
+[video_3_full_path]
+...
+```
+
+### Training
+For example, set `latent_channels=16/64` and `quant_method=hierarchical_binary_quant_round_4` in `scripts/hbq_tokenizer_train.sh`, then run:
+```bash
+cd grn/tokenizer
+bash scripts/hbq_tokenizer_train.sh
+```
+
+### Evaluation
+For example, set `latent_channels=16/64` and `quant_method=hierarchical_binary_quant_round_4` in `scripts/hbq_tokenizer_train.sh`, then run:
+```bash
+cd grn/tokenizer
+bash scripts/hbq_tokenizer_eval.sh
+```
+
+---
+
 ## 🖼️ Class-to-Image
 
-### Dataset
+### Data
 Download [ImageNet](http://image-net.org/download) dataset, and place it in your `IMAGENET_PATH`.
 
 ### Training
@@ -238,35 +275,36 @@ We use [torch-fidelity](https://github.com/LTH14/torch-fidelity) to evaluate FID
 
 ## 🎨 Text-to-Image
 ### Data
-Refer to `data/toy_data/jsonls/000001/0001_0008_000000100.jsonl`
+Refer to `data/toy_data/jsonls/000001/0001_0800_000000100.jsonl`
 ```
 {"image_path": "[image_path_1]", "long_caption": "xxx", "long_caption_type": "caption-InternVL2.0", "text": "", "short_caption_type": "blip2_caption", "width": 1080, "height": 1920}
 {"image_path": "[image_path_2]", "long_caption": "xxx", "long_caption_type": "caption-InternVL2.0", "text": "", "short_caption_type": "blip2_caption", "width": 1080, "height": 1920}
 ...
 ```
 
-### Train
-Run `bash scripts/train_GRN_ind_t2i.sh`
+### Training
+Run `bash scripts/t2iv/train_GRN_bit_t2iv.sh`
 
 ### Inference
 
-You can simply run `python3 t2i_infer.py` or use the following code:
+You can simply run `python3 tools/t2i_infer.py` or use the following code:
 
 ```python
 from PIL import Image
-from grn_pipeline import GRNPipeline
+from tools.grn_pipeline import GRNPipeline
 
 # Load pipeline
 pipeline = GRNPipeline.from_pretrained(
     hf_repo_id='bytedance-research/GRN',
     task='T2I',
     pn='1M', 
+    model='GRN2b',
     device='cpu',
 ).to('cuda')
 
 # Generate one image
 result = pipeline(
-    prompt="A cute cat playing in the garden",
+    prompt="<T2I>" + "A cute cat playing in the garden",
     guidance_scale=3.0,
     temperature=1.1,
     complexity_aware_Tmin=10,
@@ -287,29 +325,30 @@ image.save('./generated_image.jpg')
 
 ## 🎬 Text-to-Video
 ### Data
-Refer to `data/toy_data/jsonls/000001/0001_0008_000000100.jsonl`
+Refer to `data/toy_data/jsonls/000001/0001_0800_000000100.jsonl`
 ```
 {"video_path": "[video_path_1]", "begin_frame_id": xxx, "end_frame_id": xxx, "quality_prompt": "There is text in the video.", "fps": 25.0, "duration": 3.88, "width": 1280, "height": 720, "caption": [{"type": "short", "content": "[short_caption]"}, {"type": "medium", "content": "[medium_caption]"}, {"type": "long", "content": "[long_caption]"}]}
 {"video_path": "[video_path_1]", "begin_frame_id": xxx, "end_frame_id": xxx, "quality_prompt": "The quality is very high!", "fps": 25.0, "duration": 3.88, "width": 1280, "height": 720, "caption": [{"type": "short", "content": "[short_caption]"}, {"type": "medium", "content": "[medium_caption]"}, {"type": "long", "content": "[long_caption]"}]}
 ...
 ```
 
-### Train
-Run `bash scripts/train_GRN_ind_t2v.sh`
+### Training
+Run `bash scripts/t2iv/train_GRN_bit_t2iv.sh`
 
 ### Inference
 
-You can simply run `python3 t2v_infer.py` or use the following code:
+You can simply run `python3 tools/t2v_infer.py` or use the following code:
 
 ```python
-from grn_pipeline import GRNPipeline
+from tools.grn_pipeline import GRNPipeline
 
 # Load pipeline
 pipeline = GRNPipeline.from_pretrained(
-    hf_repo_id='bytedance-research/GRN', 
-    task='T2V', 
+    hf_repo_id='bytedance-research/GRN',
+    task='T2V',
     pn='0.41M', 
-    device='cpu'
+    model='GRN2b',
+    device='cpu',
 ).to('cuda')
 
 # Generate one video
@@ -333,41 +372,6 @@ video_file = result.videos[0]
 
 ---
 
-## 📦 HBQ Tokenizer
-
-### Data
-Image Dataset, e.g., data_root/username/labels/imagenet/train.txt:
-```
-[image_1_full_path]
-[image_2_full_path]
-[image_3_full_path]
-...
-```
-
-Video Dataset, e.g., data_root/username/labels_hanjian/high-quality-video/horizontal_videos.txt
-```
-[video_1_full_path]
-[video_2_full_path]
-[video_3_full_path]
-...
-```
-
-### Training
-For example, set `latent_channels=16/64` and `quant_method=hierarchical_binary_quant_round_4` in `scripts/hbq_tokenizer_train.sh`, then run:
-```bash
-cd grn/tokenizer
-bash scripts/hbq_tokenizer_train.sh
-```
-
-### Evaluation
-For example, set `latent_channels=16/64` and `quant_method=hierarchical_binary_quant_round_4` in `scripts/hbq_tokenizer_train.sh`, then run:
-```bash
-cd grn/tokenizer
-bash scripts/hbq_tokenizer_eval.sh
-```
-
----
-
 ## 📧 Contact
 
 If you are interested in scaling GRN for image generation / image editing / video generation / video editing / unified model directions, please feel free to reach out!

app.py

@@ -7,7 +7,7 @@ import spaces
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
-from grn_pipeline import GRNPipeline
+from tools.grn_pipeline import GRNPipeline
 
 # Global pipeline
 pipe = None

environment.yaml

@@ -1,19 +0,0 @@
-name: grn
-channels:
-  - pytorch
-  - defaults
-  - nvidia
-dependencies:
-  - python=3.10
-  - pip=22.3
-  - pytorch-cuda=12.4
-  - pytorch=2.5.1
-  - torchvision=0.20.1
-  - numpy=1.22
-  - pip:
-    - opencv-python==4.11.0.86
-    - timm==0.9.12
-    - tensorboard==2.10.0
-    - scipy==1.9.1
-    - einops==0.8.1
-    - gdown==5.2.0

grn/models/grn.py

@@ -606,7 +606,7 @@ class GRN(nn.Module):
             pt_unshift = (cur_inner_round_si + 1) / (args.complexity_aware_Tmax - 1)
             pt_shift = shift_pt(min(1., pt_unshift), args.snr_shift)
             next_pt = 1 - np.cos(np.pi/2*pt_shift)
-            next_pt = next_pt * 0.999
+            next_pt = next_pt * 0.95
 
             pred_cond_labels = torch.argmax(pred_cond_probs, dim=-1) # [B,thw,d]
             pred_cond_labels = bld_to_bthwd(pred_cond_labels, pt, ph, pw)

grn/tokenizer/videovae/models/hbq_tokenizer.py

@@ -16,6 +16,13 @@ from videovae.utils.dynamic_resolution_two_pyramid import dynamic_resolution_thw
 CACHE_T = 2
 
 
+def _cache_tensor(x, length=CACHE_T, detach=False):
+    cache = x[:, :, -length:, :, :]
+    if detach:
+        cache = cache.detach()
+    return cache.clone()
+
+
 class CausalConv3d(nn.Conv3d):
     """
     Causal 3d convolusion.
@@ -114,13 +121,22 @@ class Resample(nn.Module):
     def forward(self, x, feat_cache=None, feat_idx=[0]):
         b, c, t, h, w = x.size()
         if self.mode == "upsample3d":
-            if feat_cache is not None:
+            if feat_cache is None:
+                if t > 1:
+                    first_frame = x[:, :, :1, :, :]
+                    x = self.time_conv(x[:, :, 1:, :, :])
+                    x = x.reshape(b, 2, c, t - 1, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :],
+                                     x[:, 1, :, :, :, :]), 3)
+                    x = x.reshape(b, c, (t - 1) * 2, h, w)
+                    x = torch.cat([first_frame, x], 2)
+            else:
                 idx = feat_idx[0]
                 if feat_cache[idx] is None:
                     feat_cache[idx] = "Rep"
                     feat_idx[0] += 1
                 else:
-                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    cache_x = _cache_tensor(x, detach=self.training)
                     if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
                             feat_cache[idx] != "Rep"):
                         # cache last frame of last two chunk
@@ -157,13 +173,19 @@ class Resample(nn.Module):
         x = rearrange(x, "(b t) c h w -> b c t h w", t=t) # this 4 lines do spatial down / up sample
 
         if self.mode == "downsample3d":
-            if feat_cache is not None:
+            if feat_cache is None:
+                if t >= 3:
+                    x = torch.cat([x[:, :, :1, :, :], self.time_conv(x)], 2)
+                else:
+                    x = x[:, :, :1, :, :]
+            else:
                 idx = feat_idx[0]
                 if feat_cache[idx] is None:
-                    feat_cache[idx] = x.clone()
+                    feat_cache[idx] = _cache_tensor(
+                        x, length=1, detach=self.training)
                     feat_idx[0] += 1
                 else:
-                    cache_x = x[:, :, -1:, :, :].clone()
+                    cache_x = _cache_tensor(x, length=1, detach=self.training)
                     x = self.time_conv(
                         torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
                     feat_cache[idx] = cache_x
@@ -218,7 +240,7 @@ class ResidualBlock(nn.Module):
         for layer in self.residual:
             if isinstance(layer, CausalConv3d) and feat_cache is not None:
                 idx = feat_idx[0]
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                cache_x = _cache_tensor(x, detach=self.training)
                 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                     # cache last frame of last two chunk
                     cache_x = torch.cat(
@@ -255,6 +277,7 @@ class AttentionBlock(nn.Module):
         nn.init.zeros_(self.proj.weight)
 
     def forward(self, x):
+        return x
         identity = x
         b, c, t, h, w = x.size()
         x = rearrange(x, "b c t h w -> (b t) c h w")
@@ -447,7 +470,7 @@ class Down_ResidualBlock(nn.Module):
         self.downsamples = nn.Sequential(*downsamples)
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
-        x_copy = x.clone()
+        x_copy = x
         for module in self.downsamples:
             x = module(x, feat_cache, feat_idx)
 
@@ -489,7 +512,7 @@ class Up_ResidualBlock(nn.Module):
         self.upsamples = nn.Sequential(*upsamples)
 
     def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
-        x_main = x.clone()
+        x_main = x
         for module in self.upsamples:
             x_main = module(x_main, feat_cache, feat_idx)
         if self.avg_shortcut is not None:
@@ -560,7 +583,7 @@ class Encoder3d(nn.Module):
 
         if feat_cache is not None:
             idx = feat_idx[0]
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            cache_x = _cache_tensor(x, detach=self.training)
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                 cache_x = torch.cat(
                     [
@@ -594,7 +617,7 @@ class Encoder3d(nn.Module):
         for layer in self.head:
             if isinstance(layer, CausalConv3d) and feat_cache is not None:
                 idx = feat_idx[0]
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                cache_x = _cache_tensor(x, detach=self.training)
                 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                     cache_x = torch.cat(
                         [
@@ -672,7 +695,7 @@ class Decoder3d(nn.Module):
     def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
         if feat_cache is not None:
             idx = feat_idx[0]
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            cache_x = _cache_tensor(x, detach=self.training)
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                 cache_x = torch.cat(
                     [
@@ -699,13 +722,13 @@ class Decoder3d(nn.Module):
             if feat_cache is not None:
                 x = layer(x, feat_cache, feat_idx, first_chunk)
             else:
-                x = layer(x)
+                x = layer(x, first_chunk=first_chunk)
 
         ## head
         for layer in self.head:
             if isinstance(layer, CausalConv3d) and feat_cache is not None:
                 idx = feat_idx[0]
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                cache_x = _cache_tensor(x, detach=self.training)
                 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                     cache_x = torch.cat(
                         [
@@ -780,30 +803,33 @@ class WanVAE_(nn.Module):
         x_recon = self.decode(mu, scale)
         return x_recon, mu
 
-    def encode(self, x, scale=None):
+    def encode(self, x, scale=None, use_cache=None):
         self.clear_cache()
         x = patchify(x, patch_size=2)
-        # import pdb; pdb.set_trace()
-        # if self.training:
-        #     out = self.encoder(x)
-        # else:
-        t = x.shape[2]
-        iter_ = 1 + (t - 1) // 4
-        for i in range(iter_):
-            self._enc_conv_idx = [0]
-            if i == 0:
-                out = self.encoder(
-                    x[:, :, :1, :, :],
-                    feat_cache=self._enc_feat_map,
-                    feat_idx=self._enc_conv_idx,
-                )
-            else:
-                out_ = self.encoder(
-                    x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
-                    feat_cache=self._enc_feat_map,
-                    feat_idx=self._enc_conv_idx,
-                )
-                out = torch.cat([out, out_], 2)
+        if use_cache is None:
+            use_cache = not self.training
+        if use_cache:
+            t = x.shape[2]
+            iter_ = 1 + (t - 1) // 4
+            outs = []
+            for i in range(iter_):
+                self._enc_conv_idx = [0]
+                if i == 0:
+                    out_ = self.encoder(
+                        x[:, :, :1, :, :],
+                        feat_cache=self._enc_feat_map,
+                        feat_idx=self._enc_conv_idx,
+                    )
+                else:
+                    out_ = self.encoder(
+                        x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                        feat_cache=self._enc_feat_map,
+                        feat_idx=self._enc_conv_idx,
+                    )
+                outs.append(out_)
+            out = torch.cat(outs, 2) if len(outs) > 1 else outs[0]
+        else:
+            out = self.encoder(x)
         mu, log_var = self.conv1(out).chunk(2, dim=1)
         if scale is not None:
             if isinstance(scale[0], torch.Tensor):
@@ -812,15 +838,15 @@ class WanVAE_(nn.Module):
             else:
                 mu = (mu - scale[0]) * scale[1]
         self.clear_cache()
-        if self.other_args.encoder_out_type == 'feature_tanh':
+        if self.encoder_out_type == 'feature_tanh':
             mu = torch.tanh(mu)
-        elif self.other_args.encoder_out_type == 'feature':
+        elif self.encoder_out_type == 'feature':
             pass
         else:
-            raise ValueError(f'{self.other_args.encoder_out_type=} is not supported!')
+            raise ValueError(f'{self.encoder_out_type=} is not supported!')
         return mu
 
-    def decode(self, z, scale=None):
+    def decode(self, z, scale=None, use_cache=None, **kwargs):
         self.clear_cache()
         if scale is not None:
             if isinstance(scale[0], torch.Tensor):
@@ -829,18 +855,19 @@ class WanVAE_(nn.Module):
             else:
                 z = z / scale[1] + scale[0]
         x = self.conv2(z)
-        # if self.training:
-        #     out = self.decoder(x)
-        # else:
-        iter_ = z.shape[2]
-        for i in range(iter_):
-            self._conv_idx = [0]
-            if i == 0:
-                out = self.decoder(
-                    x[:, :, i:i + 1, :, :],
-                    feat_cache=self._feat_map,
-                    feat_idx=self._conv_idx,
-                    first_chunk=True,
+        if use_cache is None:
+            use_cache = not self.training
+        if use_cache:
+            iter_ = z.shape[2]
+            outs = []
+            for i in range(iter_):
+                self._conv_idx = [0]
+                if i == 0:
+                    out_ = self.decoder(
+                        x[:, :, i:i + 1, :, :],
+                        feat_cache=self._feat_map,
+                        feat_idx=self._conv_idx,
+                        first_chunk=True,
                 )
             else:
                 out_ = self.decoder(
@@ -848,7 +875,10 @@ class WanVAE_(nn.Module):
                     feat_cache=self._feat_map,
                     feat_idx=self._conv_idx,
                 )
-                out = torch.cat([out, out_], 2)
+                outs.append(out_)
+            out = torch.cat(outs, 2) if len(outs) > 1 else outs[0]
+        else:
+            out = self.decoder(x, first_chunk=True)
         out = unpatchify(out, patch_size=2)
         self.clear_cache()
         return out

requirements.txt

@@ -15,3 +15,4 @@ ftfy>=6.1.1
 transformers>=4.35.0
 regex>=2023.10.3
 pyyaml>=6.0
+flash-attn-4==4.0.0b14

grn_pipeline.py → tools/grn_pipeline.py

@@ -46,12 +46,12 @@ class GRNPipeline:
             from huggingface_hub import hf_hub_download, snapshot_download
             print(f"download weights from Hugging Face Hub: {hf_repo_id}")
             if task == 'T2I':
-                model_path = hf_hub_download(repo_id=hf_repo_id, filename="GRN_T2I_2B_FSA_94600.pth")
+                model_path = hf_hub_download(repo_id=hf_repo_id, filename="GRN_T2I_2B_FSA_137200.pth")
             elif task == 'T2V':
-                model_path = hf_hub_download(repo_id=hf_repo_id, filename="GRN_T2V_2B.pth")
+                model_path = hf_hub_download(repo_id=hf_repo_id, filename="GRN_T2V_2B_FSA_sft3800_non_ema.pth")
             else:
                 raise ValueError(f"Unknown task: {task}")
-            vae_path = hf_hub_download(repo_id=hf_repo_id, filename="HBQ_tokenizer_64dim_M4.ckpt")
+            vae_path = hf_hub_download(repo_id=hf_repo_id, filename="HBQ_image_video_tokenizer_64dim_M4_20260618.ckpt")
             snapshot_path = snapshot_download(repo_id=hf_repo_id, allow_patterns="umt5-xxl/**")
             text_encoder_ckpt = os.path.join(snapshot_path, "umt5-xxl")
             print(os.listdir(snapshot_path))

t2i_infer.py → tools/t2i_infer.py

@@ -1,5 +1,5 @@
 from PIL import Image
-from grn_pipeline import GRNPipeline
+from tools.grn_pipeline import GRNPipeline
 
 # Load pipeline
 pipeline = GRNPipeline.from_pretrained(
@@ -7,7 +7,7 @@ pipeline = GRNPipeline.from_pretrained(
     task='T2I',
     pn='1M', 
     model='GRN2b',
-    use_slow_attn=True,
+    use_slow_attn=False,
     device='cpu',
 ).to('cuda')
 

t2v_infer.py → tools/t2v_infer.py

@@ -1,4 +1,4 @@
-from grn_pipeline import GRNPipeline
+from tools.grn_pipeline import GRNPipeline
 
 negative_prompt = (
     # --- quality ---
@@ -25,16 +25,18 @@ pipeline = GRNPipeline.from_pretrained(
     hf_repo_id='bytedance-research/GRN', 
     task='T2V', 
     pn='0.41M', 
+    model='GRN2b',
+    use_slow_attn=False,
     device='cpu'
 ).to('cuda')
 
-prompt="The man, of medium build with short, dark, curly hair, stands centered in the frame, wearing a simple white t-shirt that contrasts with the greenery behind him. He holds a dark smartphone, likely a modern model with a triple-lens camera setup, in his right hand, angled slightly toward his body. His gaze is fixed on the screen, and his facial expression shifts subtly\u2014smiling, nodding, and occasionally pursing his lips\u2014as if reacting to content on the phone. The background features a mix of tall green trees and shrubs, with a light blue metal fence running horizontally across the mid-ground, suggesting a garden or rural boundary. The overcast sky diffuses the light, creating soft shadows and a calm, neutral atmosphere. The man\u2019s slight head movements and micro-expressions indicate engagement, possibly reading or responding to a message or video. The composition places him as the focal point, with the natural, slightly blurred background reinforcing his isolation in the moment. The relative stillness of the scene, apart from his subtle gestures, suggests a private, introspective interaction with technology in a serene outdoor setting"
-
+prompt="The video captures a male performer on stage, wearing a black cap, black t-shirt, and a black beaded bracelet on his left wrist, with a tattoo visible on his left forearm. He holds a microphone close to his mouth with his left hand while raising his right arm in a dynamic gesture, suggesting energetic performance. The stage is illuminated with intense blue and purple lighting, creating a moody atmosphere; a focused spotlight beam is visible in the background, adding depth. The performer’s facial expression is intense, eyes closed or squinting, indicating emotional engagement. The camera maintains a close-up, slightly angled shot of his upper body, with minimal movement, emphasizing his actions and expressions. Across the frames, the lighting subtly shifts, enhancing the visual dynamics of the performance without altering the scene’s core composition"
+prompt='一个头戴耳机的男人正在一个工作室里对着麦克风唱歌，偶尔晃动头部'
 # Generate one video
 result = pipeline(
-    prompt=f"{prompt}. masterpiece, high quality.",
+    prompt=f"<T2V>{prompt}. The quality is very high!",
     negative_prompt=negative_prompt,
-    guidance_scale=4.0,
+    guidance_scale=3.0,
     temperature=1.0,
     complexity_aware_Tmin=10,
     complexity_aware_Tmax=50,