feat: Add model files with LFS

.gitattributes

@@ -0,0 +1,4 @@
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+results/kling_9k_9k/best_acc.pth
+models/clip/bpe_simple_vocab_16e6.txt.gz

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

Preprocess/compress.py

@@ -0,0 +1,38 @@
+import subprocess
+from glob import glob
+
+def convert_and_compress_video(input_video_path, output_video_path, crf=23, preset='medium'):
+    """
+    使用ffmpeg转换并压缩视频
+    :param input_video_path: 输入视频文件路径
+    :param output_video_path: 输出压缩后视频文件路径
+    :param crf: Constant Rate Factor，值越小质量越好，但文件也越大，默认为23是一个平衡点
+    :param preset: 预设值，影响压缩速度和文件大小，如'ultrafast', 'fast', 'medium', 'slow', 'veryslow'等，默认'medium'
+    """
+    # 检查文件是否为GIF
+    if input_video_path.lower().endswith('.gif'):
+        # 构建ffmpeg命令
+        # 首先将GIF转换为MP4
+        convert_cmd = f'ffmpeg -i "{input_video_path}" -vf "palettegen" -y palette.png'
+        subprocess.run(convert_cmd, shell=True, check=True)
+        convert_cmd = f'ffmpeg -i "{input_video_path}" -i palette.png -lavfi "paletteuse" -c:v libx264 -preset {preset} -crf {crf} -c:a copy "{output_video_path}"'
+    else:
+        # 直接压缩视频
+        convert_cmd = f'ffmpeg -i "{input_video_path}" -c:v libx264 -preset {preset} -crf {crf} -c:a copy "{output_video_path}"'
+    
+    # 执行ffmpeg命令
+    try:
+        subprocess.run(convert_cmd, shell=True, check=True)
+        print(f"视频压缩完成，输出文件：{output_video_path}")
+    except subprocess.CalledProcessError as e:
+        print(f"视频压缩失败：{e}")
+
+video_paths = '/home3/Transformer'
+all_dirs = glob(video_paths+'/*')
+output_dirs = '/home3/robust/compress/Transformer_'
+
+for path in all_dirs:
+    name = path.split('/')[-1]
+    out_path = output_dirs + name.replace('.gif', '.mp4') if path.lower().endswith('.gif') else output_dirs + name
+    print(name)
+    convert_and_compress_video(path, out_path, crf=28)

Preprocess/folder2csv.py

@@ -0,0 +1,72 @@
+import os
+import csv
+import pandas as pd
+from pandas import Series, DataFrame
+from glob import glob
+import os
+
+def count_images_in_folder(folder_path):
+    image_count = 0
+    image_names = []  
+    for file_name in os.listdir(folder_path):
+        if file_name.endswith('.png') or file_name.endswith('.jpg') or file_name.endswith('.jpeg'):
+            image_count += 1
+            image_names.append(int(file_name.split('.')[0]))
+    image_names.sort()
+    return image_count, image_names
+
+
+folder_path = './SD_frames'
+all_dirs = []
+
+for root, dirs, files in os.walk(folder_path):
+    for dir in dirs:
+        all_dirs.append(os.path.join(root, dir))
+
+label = list()
+save_path = list()
+frame_counts = list()
+frame_seq_counts = list()
+content_paths = list()
+chinese_labels = list()
+
+
+for video_path in all_dirs:
+    frame_paths = glob(video_path + '/*')
+    temp_frame_count, temp_frame_seqs = count_images_in_folder(video_path)
+    if temp_frame_count == 0:
+        continue
+
+    for frame in frame_paths:
+        content_path = frame.split('/')[1:-1]
+        content_path = '/'.join(content_path)
+        # input your own path
+        content_path = '/home/AIGC_Video_Det/SD/' + content_path
+      
+        frame_path = frame.split('/')[1:]
+        frame_path = '/'.join(frame_path)
+        frame_path = '/home/AIGC_Video_Det/SD/' + frame_path
+
+        print(content_path, frame_path)
+        label.append(str(1))
+        frame_counts.append(int(temp_frame_count))
+        frame_seq_counts.append(temp_frame_seqs)
+        save_path.append(frame_path)
+        content_paths.append(content_path)
+        chinese_labels.append('AIGC视频')
+        # chinese_labels.append('真实视频')
+        break
+
+dic={
+    'content_path': Series(data=content_paths),
+    'image_path': Series(data=save_path),
+    'type_id': Series(data=chinese_labels),
+    'label': Series(data=label),
+    'frame_len': Series(data=frame_counts),
+    'frame_seq': Series(data=frame_seq_counts)
+}
+
+print(dic)
+pd.DataFrame(dic).to_csv('SD.csv', encoding='utf-8', index=False)
+
+

Preprocess/video2frame.py

@@ -0,0 +1,57 @@
+import os
+from glob import glob
+from moviepy.editor import VideoFileClip
+import multiprocessing
+import cv2, math
+
+def get_video_length(file_path):
+    video = VideoFileClip(file_path)
+    return video.duration
+
+def process_video(video_path):
+    video_name = video_path.split('/')[-1]
+    video_name = video_name.split('.')[:-1]
+    video_name = '.'.join(video_name)
+
+    path = video_path.split('/')[1:-1]
+    path = '/'.join(path)
+
+    image_path = './SD_frames/'+path+'/'+ video_name+'/'
+    print(video_name)
+    if os.path.exists(image_path):
+        print("路径存在")
+    else:
+        print(video_name, "路径不存在")
+        try:
+            try:
+                video_length = get_video_length(video_path)
+                print(video_name, f"视频长度为：{video_length} 秒")
+                os.makedirs(os.path.dirname(image_path), exist_ok=True)
+
+                if video_length >= 4 :
+                    inter_val = 2
+                    os.system(f"cd {image_path} | ffmpeg -loglevel quiet -i {video_path} -r {inter_val} {image_path}%d.jpg")
+                else:
+                    inter_val = math.ceil(8 / video_length)
+                    os.system(f"cd {image_path} | ffmpeg -loglevel quiet -i {video_path} -r {inter_val} {image_path}%d.jpg")
+ 
+            except Exception as e:
+                print("发生异常：", str(e))
+        except:
+            print("Skip")
+
+if __name__ == '__main__':
+    print("Getting frames!!")
+    video_paths = './SD'
+    all_dirs = []
+    all_dirs = glob(video_paths+'/*')
+    
+    print(all_dirs)
+    
+    pool = multiprocessing.Pool(processes=8)
+    pool.map(process_video, all_dirs)
+    pool.close()
+    pool.join()
+
+
+

README.md

@@ -0,0 +1,88 @@
+This is the official code of paper 'DeMamba: AI-Generated Video Detection on Million-Scale GenVideo Benchmark'.
+
+## :boom: News!!! 
+- (25.09.24) We have released a lightweight version of GenVideo, named [GenVideo-100K](https://modelscope.cn/datasets/cccnju/GenVideo-100K), with 10,000 samples for each forgery category.
+- (24.12.11) We are pleased to announce that our AI-generated content (AIGC) video detection model, developed based on GenVideo, has passed the evaluation by the China Academy of Information and Communications Technology (CAICT) and achieved the highest rating, making us the first organization in China to be officially registered and approved.[[certificate](https://github.com/chenhaoxing/DeMamba/blob/main/figs/xty.jpg)][[report](https://mp.weixin.qq.com/s/OoW7EI1QoSrQ3FIftfbudg)]
+
+
+## :dart: Todo
+
+- [x] Release source code.
+- [x] Release dataset. 
+
+
+## :file_folder: Dataset download
+![](figs/tab_fig.jpg)
+
+### Data preparation process
+ - Download the original videos.
+   
+   - Generated videos: all generated videos can download at [ModelScope](https://modelscope.cn/collections/Gen-Video-7de46cd6846f4e).
+     
+   - Real videos: The data from the MSRVTT dataset is contained within the GenVideo-Val.zip file. We also provided the selected Youku videos in previous link . For Kinetics-400, you will need to download it yourself at [https://github.com/cvdfoundation/kinetics-dataset](https://github.com/cvdfoundation/kinetics-dataset).
+     
+ - Preprocess the video and get the data list csv file.
+
+Statistics of real and generated videos in the GenVideo dataset:
+| **Video Source**                                    | **Type** | **Task** | **Time** | **Resolution** | **FPS** | **Length** | **Training Set** | **Testing Set** |
+|-----------------------------------------------------|----------|----------|----------|----------------|---------|------------|------------------|----------------|
+| Kinetics-400               | Real     | -        | 17.05    | 224-340        | -       | 5-10s      | 260,232          | -              |
+| Youku-mPLUG                     | Real     | -        | 23.07    | -              | -       | 10-120s    | 953,279          | -             |
+| MSR-VTT                           | Real     | -        | 16.05    | -              | -       | 10-30s     | -                | 10,000         |
+| ZeroScope                       | Fake     | T2V      | 23.07    | 1024×576       | 8       | 3s         | 133,169          | -             |
+| I2VGen-XL                         | Fake     | I2V      | 23.12    | 1280×720       | 8       | 2s         | 61,975           | -            |
+| SVD                     | Fake     | I2V      | 23.12    | 1024×576       | 8       | 4s         | 149,026          | -              | 
+| VideoCrafte          | Fake     | T2V      | 24.01    | 1024×576       | 8       | 2s         | 39,485           | -              | 
+| Pika                                   | Fake     | T2V&I2V  | 24.02    | 1088×640       | 24      | 3s         | 98,377           |        |
+| DynamiCrafter         | Fake     | I2V      | 24.03    | 1024×576       | 8       | 3s         | 46,205          | -              | 
+| SD                             | Fake     | T2V&I2V  | 23-24    | 512-1024       | 8       | 2-6s       | 200,720          | -              |                 
+| SEINE                         | Fake     | I2V      | 24.04    | 1024×576       | 8       | 2-4s       | 24,737            | -              |                 
+| Latte                           | Fake     | T2V      | 24.03    | 512×512        | 8       | 2s         | 149,979          | -              |                 
+| OpenSora                           | Fake     | T2V      | 24.03    | 512×512        | 8       | 2s         | 177,410          | -              |                 
+| ModelScope               | Fake     | T2V      | 23.03    | 256×256        | 8       | 4s         | -                | 700            |
+| MorphStudio                     | Fake     | T2V      | 23.08    | 1280×720       | 8       | 2s         | -                | 700            |                 
+| MoonValley                       | Fake     | T2V      | 24.01    | 1024×576       | 16      | 3s         | -                | 626            |                 
+| HotShot                           | Fake     | T2V      | 23.10    | 672×384        | 8       | 1s         | -                | 700            |                 
+| Show_1                       | Fake     | T2V      | 23.10    | 576×320        | 8       | 4s         | -                | 700            |                 
+| Gen2                     | Fake     | I2V&T2V  | 23.09    | 896×512        | 24      | 4s         | -                | 1,380          |                 
+| Crafter               | Fake     | T2V      | 23.04    | 256×256        | 8       | 4s         | -                | 1,400          |                 
+| Lavie                                 | Fake     | T2V      | 23.09    | 1280×2048      | 8       | 2s         | -                | 1,400          |                 
+| Sora                                 | Fake     | T2V      | 24.02    | -              | -       | -60s       | -                | 56             |                
+| WildScrape                                          | Fake     | T2V&I2V  | 24       | 512-1024       | 8-16    | 2-6s       | -                | 926            |
+| **Total Count**                                     | -        | -        | -        | -              | -       | -          | 2,294,594       | 19,588         | 
+
+## :snake: Detail Mamba (DeMamba)
+
+![](figs/logo.png)
+<p align="center"><em>In memory of Kobe Bryant (generated by GPT-4o)</em></p>
+
+> "Determination wins games, but Detail wins championships." — *Kobe Bryant, in his Show Detail, 2018*
+
+![](figs/VFOD.png)
+<p align="center"><em>The overall framework of our Detail Mamba (DeMamba)</em></p>
+
+
+## :space_invader: Citing GenVideo&DeMamba
+If you use GenVideo or DeMamba in your research or use the codebase here, please use the following BibTeX entry.
+
+```BibTeX
+@article{DeMamba,
+      title={DeMamba: AI-Generated Video Detection on Million-Scale GenVideo Benchmark},
+      author={Haoxing Chen and Yan Hong and Zizheng Huang and Zhuoer Xu and Zhangxuan Gu and Yaohui Li and Jun Lan and Huijia Zhu and Jianfu Zhang and Weiqiang Wang and Huaxiong Li},
+      journal={arXiv preprint arXiv:2405.19707},
+      year={2024}
+}
+```
+
+## Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=chenhaoxing/DeMamba&type=Date)](https://star-history.com/#chenhaoxing/DeMamba&Date)
+
+## Acknowledgement
+Many thanks to the nice work of [STIL](https://github.com/wizyoung/STIL-DeepFake-Video-Detection), [CLIP](https://github.com/openai/CLIP), [XCLIP](https://github.com/microsoft/VideoX/tree/master/X-CLIP), [NPR](https://github.com/chuangchuangtan/NPR-DeepfakeDetection/tree/main) and [VideoMAE](https://github.com/MCG-NJU/VideoMAE-Action-Detection). 
+
+## :email: Contact
+If you have any questions, feel free to contact us: hx.chen@hotmail.com.
+
+
+

commands.md

@@ -0,0 +1,3 @@
+python train.py --config /home/kalpit/workspace/aigc/repos/DeMamba/configs/my_config.yaml
+
+python train.py --config /home/kalpit/workspace/aigc/repos/DeMamba/configs/vid_grand_v1_train.yaml

create_csv.py

@@ -0,0 +1,44 @@
+import os
+import csv
+
+# Source directories
+real_dir = "/home/kalpit/workspace/aigc/data/ShareVeo3/test/0_real"
+fake_dir = "/home/kalpit/workspace/aigc/data/ShareVeo3/test/1_fake"
+
+# Output CSV path
+output_csv = "/home/kalpit/workspace/aigc/repos/DeMamba/veo_test.csv"
+
+# Function to get all image paths from a directory
+def get_image_paths(directory, label):
+    image_paths = []
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
+                full_path = os.path.join(root, file)
+                image_paths.append({
+                    "content_path": full_path,
+                    "frame_seq": [full_path],  # list containing the single frame path
+                    "label": label
+                })
+    return image_paths
+
+# Collect all images
+data = []
+data.extend(get_image_paths(real_dir, 0))
+data.extend(get_image_paths(fake_dir, 1))
+
+# Write to CSV
+with open(output_csv, 'w', newline='', encoding='utf-8') as csvfile:
+    fieldnames = ["content_path", "frame_seq", "label"]
+    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+
+    writer.writeheader()
+    for row in data:
+        # Convert frame_seq list to string as shown in your example
+        writer.writerow({
+            "content_path": row["content_path"],
+            "frame_seq": str(row["frame_seq"]),
+            "label": row["label"]
+        })
+
+print(f"CSV saved at: {output_csv}")

create_submission.py

@@ -0,0 +1,35 @@
+import pandas as pd
+import re
+
+# Input and output files
+input_csv = "results.csv"
+output_csv = "submission.csv"
+
+# Load frame-level results
+df = pd.read_csv(input_csv)
+
+# Extract base video ID (remove _f###.jpg)
+def extract_video_id(filename):
+    return re.sub(r"_f\d+\.\w+$", "", filename)
+
+df["video_id"] = df["file_name"].apply(extract_video_id)
+
+# Aggregate probabilities per video (mean of frame scores)
+video_scores = (
+    df.groupby("video_id")["predicted_prob"]
+    .mean()
+    .reset_index()
+    .rename(columns={"video_id": "id", "predicted_prob": "score"})
+)
+
+# Assign label: generated if >0.5 else real
+video_scores["pred"] = video_scores["score"].apply(
+    lambda x: "generated" if x > 0.5 else "real"
+)
+
+# Reorder columns
+video_scores = video_scores[["id", "pred", "score"]]
+
+# Save submission file
+video_scores.to_csv(output_csv, index=False)
+print(f"Saved submission file to {output_csv}")

dataloader.py

@@ -0,0 +1,281 @@
+import torch.utils.data as data
+from torch.utils.data import Dataset, DataLoader
+import pandas as pd
+import torch
+import albumentations
+import random
+import os
+import numpy as np
+import cv2
+import math
+import warnings
+
+
+def crop_center_by_percentage(image, percentage):
+    height, width = image.shape[:2]
+
+    if width > height:
+        left_pixels = int(width * percentage)
+        right_pixels = int(width * percentage)
+        start_x = left_pixels
+        end_x = width - right_pixels
+        cropped_image = image[:, start_x:end_x]
+    else:
+        up_pixels = int(height * percentage)
+        down_pixels = int(height * percentage)
+        start_y = up_pixels
+        end_y = height - down_pixels
+        cropped_image = image[start_y:end_y, :]
+
+    return cropped_image
+
+class Ours_Dataset_train(Dataset):
+    def __init__(self, index_list=None, df=None):
+        self.index_list = index_list
+        self.df = df
+        self.positive_indices = df[df['label'] == 1].index.tolist()
+        self.negative_indices = df[df['label'] == 0].index.tolist()
+        self.balanced_indices = []
+        self.resample()
+
+    def resample(self):
+        # Ensure each epoch uses a balanced dataset
+        min_samples = min(len(self.positive_indices), len(self.negative_indices))
+        self.balanced_indices.clear()
+        self.balanced_indices.extend(random.sample(self.positive_indices, min_samples))
+        self.balanced_indices.extend(random.sample(self.negative_indices, min_samples))
+        random.shuffle(self.balanced_indices)  # Shuffle to mix positive and negative samples
+
+    def __getitem__(self, idx):
+        real_idx = self.balanced_indices[idx]
+        row = self.df.iloc[real_idx]
+        video_id = row['content_path']
+        label = row['label']
+        frame_list = eval(row['frame_seq'])
+        label_onehot = [0]*2
+        select_frame_nums = 8
+
+        aug_list  = [
+                    albumentations.Resize(224, 224)
+                    ]
+
+        if random.random() < 0.5:
+            aug_list.append(albumentations.HorizontalFlip(p=1.0))
+        if random.random() < 0.5:
+            quality_score = random.randint(50, 100)
+            aug_list.append(albumentations.ImageCompression(quality_lower=quality_score, quality_upper=quality_score))
+        if random.random() < 0.3:
+            aug_list.append(albumentations.GaussNoise(p=1.0))
+        if random.random() < 0.3:
+            aug_list.append(albumentations.GaussianBlur(blur_limit=(3, 5), p=1.0))
+        if random.random() < 0.001:
+            aug_list.append(albumentations.ToGray(p=1.0))
+            
+        aug_list.append(albumentations.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0))
+        trans = albumentations.Compose(aug_list)
+
+        if len(frame_list) >= select_frame_nums:
+            start_frame = random.randint(0, len(frame_list)-select_frame_nums)
+            select_frames = frame_list[start_frame:start_frame+select_frame_nums]
+            frames = []
+            for x in frame_list[start_frame:start_frame+select_frame_nums]:
+                while True:
+                    try:
+                        temp_image_path = video_id+'/'+str(x)+'.jpg'
+                        image = download_oss_file('GenVideo/'+ temp_image_path)
+                        if video_id.startswith("real/youku"):
+                            image = crop_center_by_percentage(image, 0.15)
+                        break
+                    except Exception as e:
+                        if x+1 < len(frame_list):
+                            x = x + 1
+                        elif x - 1 >=0 :
+                            x = x - 1
+                augmented = trans(image=image)
+                image = augmented["image"]
+                frames.append(image.transpose(2,0,1)[np.newaxis,:])
+        else:
+            pad_num = select_frame_nums-len(frame_list)
+            frames = []
+            for x in frame_list:
+                temp_image_path = video_id+'/'+str(x)+'.jpg'
+                image = download_oss_file('GenVideo/'+temp_image_path)
+                if video_id.startswith("real/youku"):
+                    image = crop_center_by_percentage(image, 0.15)
+                augmented = trans(image=image)
+                image = augmented["image"]
+                frames.append(image.transpose(2,0,1)[np.newaxis,:])    
+            for i in range(pad_num):
+                frames.append(np.zeros((224,224,3)).transpose(2,0,1)[np.newaxis,:])
+        
+        label_onehot[int(label)] = 1
+        frames = np.concatenate(frames, 0)
+        frames = torch.tensor(frames[np.newaxis,:])
+        label_onehot = torch.FloatTensor(label_onehot)
+        binary_label = torch.FloatTensor([int(label)])
+
+        return self.index_list[idx], frames, label_onehot, binary_label
+
+    def __len__(self):
+        return len(self.balanced_indices)
+
+
+class Ours_Dataset_val(data.Dataset):
+    def __init__(self, cfg, index_list=None, df=None):
+        self.index_list = index_list
+        self.cfg = cfg
+        self.df = df
+        self.frame_dir = df['image_path'].tolist()
+
+    def __getitem__(self, idx):
+        aug_list  = [
+                    albumentations.Resize(224, 224),
+                    ]
+        
+        if self.cfg['task'] == 'JPEG_Compress_Attack':
+            aug_list.append(albumentations.JpegCompression(quality_lower=35, quality_upper=35,p=1.0))
+        if self.cfg['task'] == 'FLIP_Attack':
+            if random.random() < 0.5:
+                aug_list.append(albumentations.HorizontalFlip(p=1.0))
+            else:
+                aug_list.append(albumentations.VerticalFlip(p=1.0))
+        if self.cfg['task'] == 'CROP_Attack':
+            random_crop_x = random.randint(0, 16)  
+            random_crop_y = random.randint(0, 16)  
+            crop_width = random.randint(160, 208) 
+            crop_height = random.randint(160, 208)
+            aug_list.append(albumentations.Crop(x_min=random_crop_x, y_min=random_crop_y, x_max=random_crop_x+crop_width, y_max=random_crop_y+crop_height))
+            aug_list.append(albumentations.Resize(224, 224))
+
+        if self.cfg['task'] == 'Color_Attack':
+            index = random.choice([i for i in range(4)])
+            dicts = {0:[0.5,0,0,0],1:[0,0.5,0,0],2:[0,0,0.5,0],3:[0,0,0,0.5]}
+            brightness,contrast,saturation,hue = dicts[index]
+            aug_list.append(albumentations.ColorJitter(
+                brightness=brightness, contrast=contrast, saturation=saturation, hue=hue))
+
+        if self.cfg['task'] == 'Gaussian_Attack':     
+            aug_list.append(albumentations.GaussianBlur(blur_limit=(7, 7), p=1.0))
+
+        aug_list.append(albumentations.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0))
+        trans = albumentations.Compose(aug_list)
+
+
+        df_v = self.df.loc[self.index_list[idx]]
+        video_id = df_v['content_path']
+        activity_id = df_v['activity_id']
+        label = df_v['label']
+        label_onehot = [0]*2
+        frame_list = eval(df_v['frame_seq'])
+
+        select_frame_nums = 8
+
+        if len(frame_list) >= select_frame_nums:
+            start_frame = random.randint(0, len(frame_list)-select_frame_nums)
+            select_frames = frame_list[start_frame:start_frame+select_frame_nums]
+            frames = []
+            for x in frame_list[start_frame:start_frame+select_frame_nums]:
+                while True:
+                    try:
+                        temp_image_path = video_id+'/'+str(x)+'.jpg'
+                        image = download_oss_file('GenVideo/'+ temp_image_path)
+                        image = crop_center_by_percentage(image, 0.1)
+                        break
+                    except Exception as e:
+                        if x+1 < len(frame_list):
+                            x = x + 1
+                        elif x - 1 >=0 :
+                            x = x - 1
+                augmented = trans(image=image)
+                image = augmented["image"]
+                frames.append(image.transpose(2,0,1)[np.newaxis,:])
+        else:
+            pad_num = select_frame_nums-len(frame_list)
+            frames = []
+            for x in frame_list:
+                temp_image_path = video_id+'/'+str(x)+'.jpg'
+                image = download_oss_file('GenVideo/'+temp_image_path)
+                image = crop_center_by_percentage(image, 0.1)
+                augmented = trans(image=image)
+                image = augmented["image"]
+                frames.append(image.transpose(2,0,1)[np.newaxis,:])    
+            for i in range(pad_num):
+                frames.append(np.zeros((224,224,3)).transpose(2,0,1)[np.newaxis,:])
+
+        label_onehot[int(label)] = 1
+        frames = np.concatenate(frames, 0)
+        frames = torch.tensor(frames[np.newaxis,:])
+        label_onehot = torch.FloatTensor(label_onehot)
+        binary_label = torch.FloatTensor([int(label)])
+        return self.index_list[idx], frames, label_onehot, binary_label, video_id
+
+    
+    def __len__(self):
+        return len(self.index_list)
+
+
+
+def generate_dataset_loader(cfg):
+    df_train = pd.read_csv('GenVideo/datasets/train.csv')
+
+    if cfg['task'] == 'normal':
+        df_val = pd.read_csv('GenVideo/datasets/val_id.csv')
+    elif cfg['task'] == 'robust_compress':
+        df_val = pd.read_csv('GenVideo/datasets/com_28.csv')
+    elif cfg['task'] == 'Image_Water_Attack':
+        df_val = pd.read_csv('GenVideo/datasets/imgwater.csv')
+    elif cfg['task'] == 'Text_Water_Attack':
+        df_val = pd.read_csv('GenVideo/datasets/textwater.csv')
+    elif cfg['task'] == 'one2many':
+        df_val = pd.read_csv('GenVideo/datasets/val_ood.csv')
+        if cfg['train_sub_set'] == 'pika':
+            prefixes = ["fake/pika", "real"]
+            video_condition = df_train['content_path'].str.startswith(prefixes[0])
+            for prefix in prefixes[1:]:
+                video_condition |= df_train['content_path'].str.startswith(prefix)
+            df_train = df_train[video_condition]
+        elif cfg['train_sub_set'] == 'SEINE':
+            prefixes = ["fake/SEINE", "real"]
+            video_condition = df_train['content_path'].str.startswith(prefixes[0])
+            for prefix in prefixes[1:]:
+                video_condition |= df_train['content_path'].str.startswith(prefix)
+            df_train = df_train[video_condition]
+        elif cfg['train_sub_set'] == 'OpenSora':
+            prefixes = ["fake/OpenSora", "real"]
+            video_condition = df_train['content_path'].str.startswith(prefixes[0])
+            for prefix in prefixes[1:]:
+                video_condition |= df_train['content_path'].str.startswith(prefix)
+            df_train = df_train[video_condition]
+        elif cfg['train_sub_set'] == 'Latte':
+            prefixes = ["fake/Latte", "real"]
+            video_condition = df_train['content_path'].str.startswith(prefixes[0])
+            for prefix in prefixes[1:]:
+                video_condition |= df_train['content_path'].str.startswith(prefix)
+            df_train = df_train[video_condition]
+    else:
+        df_val = pd.read_csv('GenVideo/datasets/val_ood.csv')
+
+    df_train.reset_index(drop=True, inplace=True)
+    df_val.reset_index(drop=True, inplace=True)
+    
+    index_val = df_val.index.tolist()
+    index_val = index_val[:]
+
+    val_dataset = Ours_Dataset_val(cfg, index_val, df_val)
+    val_loader = torch.utils.data.DataLoader(
+            val_dataset, batch_size=cfg['val_batch_size'], shuffle=False, num_workers=cfg['num_workers'], pin_memory=True, drop_last=False
+        )
+
+    index_train = df_train.index.tolist()
+    index_train = index_train[:]
+    train_dataset = Ours_Dataset_train(index_train, df_train)
+    train_loader = torch.utils.data.DataLoader(
+            train_dataset, batch_size=cfg['train_batch_size'], shuffle=True, num_workers=cfg['num_workers'], pin_memory=True, drop_last=True
+        )
+
+    print("******* Training Video IDs", str(len(index_train))," Training Batch size ", str(cfg['train_batch_size'])," *******")
+    print("******* Testing Video IDs", str(len(index_val)), " Testing Batch size ", str(cfg['val_batch_size'])," *******")
+
+    return train_loader, val_loader
+
+

dataloader2.py

@@ -0,0 +1,246 @@
+# new_dataloader.py
+
+import torch
+import cv2
+import os
+import random
+import pandas as pd
+import albumentations
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+
+# =========== TRANSFORMATION HELPERS ===========
+
+def get_train_transforms():
+    """Defines the probabilistic augmentations for training."""
+    return albumentations.Compose([
+        albumentations.Resize(224, 224),
+        albumentations.HorizontalFlip(p=0.5),
+        albumentations.ImageCompression(quality_lower=50, quality_upper=100, p=0.5),
+        albumentations.GaussNoise(p=0.3),
+        albumentations.GaussianBlur(blur_limit=(3, 5), p=0.3),
+        albumentations.ToGray(p=0.01),
+        albumentations.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0)
+    ])
+
+def get_val_transforms(cfg):
+    """Defines augmentations for validation, handling different attack tasks from the config."""
+    aug_list = [albumentations.Resize(224, 224)]
+    
+    task = cfg.get('task', 'normal') # Use .get for safety
+
+    if task == 'JPEG_Compress_Attack':
+        aug_list.append(albumentations.JpegCompression(quality_lower=35, quality_upper=35, p=1.0))
+    elif task == 'FLIP_Attack':
+        aug_list.append(albumentations.HorizontalFlip(p=0.5)) # Original had random choice, 50% HFlip is common
+    elif task == 'CROP_Attack':
+        aug_list.append(albumentations.RandomCrop(height=192, width=192, p=1.0))
+        aug_list.append(albumentations.Resize(224, 224))
+    elif task == 'Color_Attack':
+        aug_list.append(albumentations.ColorJitter(p=1.0))
+    elif task == 'Gaussian_Attack':
+        aug_list.append(albumentations.GaussianBlur(blur_limit=(7, 7), p=1.0))
+
+    aug_list.append(albumentations.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0))
+    return albumentations.Compose(aug_list)
+
+# =========== TRAINING DATASET ===========
+
+class VideoDataset(Dataset):
+    """
+    A PyTorch Dataset for loading video frame sequences based on a DataFrame.
+    Handles class balancing for each epoch.
+    """
+    def __init__(self, df, index_list, base_data_path, transform=None, select_frame_nums=8):
+        self.df = df
+        self.index_list = index_list
+        self.base_data_path = base_data_path
+        self.transform = transform
+        self.select_frame_nums = select_frame_nums
+
+        self.positive_indices = self.df[self.df['label'] == 1].index.tolist()
+        self.negative_indices = self.df[self.df['label'] == 0].index.tolist()
+        
+        self.balanced_indices = []
+        self.resample()
+
+    def resample(self):
+        min_samples = min(len(self.positive_indices), len(self.negative_indices))
+        self.balanced_indices.clear()
+        self.balanced_indices.extend(random.sample(self.positive_indices, min_samples))
+        self.balanced_indices.extend(random.sample(self.negative_indices, min_samples))
+        random.shuffle(self.balanced_indices)
+
+    def __len__(self):
+        return len(self.balanced_indices)
+
+    def __getitem__(self, idx):
+        real_idx = self.balanced_indices[idx]
+        row = self.df.iloc[real_idx]
+        
+        video_id = row['content_path']
+        label = int(row['label'])
+        frame_list = eval(row['frame_seq'])
+
+        frames = []
+        
+        if len(frame_list) >= self.select_frame_nums:
+            start_index = random.randint(0, len(frame_list) - self.select_frame_nums)
+            selected_frames = frame_list[start_index : start_index + self.select_frame_nums]
+        else:
+            selected_frames = frame_list
+        
+        for frame_path in selected_frames:
+            try:
+                image = cv2.imread(frame_path)
+                if image is None:
+                    raise ValueError("Failed to load")
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            except Exception:
+                image = np.zeros((224, 224, 3), dtype=np.uint8)
+
+            if self.transform:
+                image = self.transform(image=image)['image']
+
+            frames.append(image.transpose(2, 0, 1)[np.newaxis, :])
+
+
+        pad_num = self.select_frame_nums - len(frames)
+        if pad_num > 0:
+            for _ in range(pad_num):
+                frames.append(np.zeros((1, 3, 224, 224)))
+
+        frames_tensor = np.concatenate(frames, axis=0)
+        frames_tensor = torch.from_numpy(frames_tensor).float().unsqueeze(0)
+
+        label_onehot = torch.zeros(2)
+        label_onehot[label] = 1.0
+        binary_label = torch.FloatTensor([label])
+        
+        original_index = self.index_list[idx]
+        return original_index, frames_tensor, label_onehot, binary_label
+
+# =========== VALIDATION DATASET ===========
+
+class VideoDatasetVal(Dataset):
+    """A compatible validation dataset loader."""
+    def __init__(self, df, index_list, base_data_path, transform=None, select_frame_nums=8):
+        self.df = df
+        self.index_list = index_list
+        self.base_data_path = base_data_path
+        self.transform = transform
+        self.select_frame_nums = select_frame_nums
+
+    def __len__(self):
+        return len(self.index_list)
+
+    def __getitem__(self, idx):
+        # Validation does not use balanced sampling, it uses the provided index directly
+        real_idx = self.index_list[idx]
+        row = self.df.iloc[real_idx]
+        
+        video_id = row['content_path']
+        label = int(row['label'])
+        frame_list = eval(row['frame_seq'])
+        
+        # This part is identical to the training dataset's __getitem__
+        frames = []
+        if len(frame_list) >= self.select_frame_nums:
+            start_index = random.randint(0, len(frame_list) - self.select_frame_nums)
+            selected_frames = frame_list[start_index : start_index + self.select_frame_nums]
+        else:
+            selected_frames = frame_list
+        
+        for frame_path in selected_frames:
+            try:
+                image = cv2.imread(frame_path)
+                if image is None:
+                    raise ValueError("Failed to load")
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            except Exception:
+                image = np.zeros((224, 224, 3), dtype=np.uint8)
+
+            if self.transform:
+                image = self.transform(image=image)['image']
+
+            frames.append(image.transpose(2, 0, 1)[np.newaxis, :])
+
+
+        pad_num = self.select_frame_nums - len(frames)
+        if pad_num > 0:
+            for _ in range(pad_num):
+                frames.append(np.zeros((1, 3, 224, 224)))
+
+        frames_tensor = np.concatenate(frames, axis=0)
+        frames_tensor = torch.from_numpy(frames_tensor).float().unsqueeze(0)
+
+        label_onehot = torch.zeros(2)
+        label_onehot[label] = 1.0
+        binary_label = torch.FloatTensor([label])
+        
+        # The original validation loader returned video_id at the end
+        return self.index_list[idx], frames_tensor, label_onehot, binary_label, video_id
+
+# =========== DATALOADER GENERATOR FUNCTION ===========
+
+def generate_dataset_loader(cfg):
+    """
+    The main function to create train and validation dataloaders using the new classes.
+    """
+    df_train = pd.read_csv('/home/kalpit/workspace/aigc/repos/DeMamba/csv/veo_train.csv')
+
+    # This logic for selecting different validation sets is preserved
+    task = cfg.get('task', 'normal')
+    if task == 'normal':
+        df_val = pd.read_csv('GenVideo/datasets/val_id.csv')
+    elif task == 'robust_compress':
+        df_val = pd.read_csv('GenVideo/datasets/com_28.csv')
+    # ... (add other elif conditions from your original script if needed) ...
+    else:
+        df_val = pd.read_csv('/home/kalpit/workspace/aigc/repos/DeMamba/csv/veo_test.csv')
+    
+    # This logic for subsetting the training data is also preserved
+    if cfg.get('train_sub_set'):
+        prefixes = [f"fake/{cfg['train_sub_set']}", "real"]
+        condition = df_train['content_path'].str.startswith(tuple(prefixes))
+        df_train = df_train[condition]
+
+    df_train.reset_index(drop=True, inplace=True)
+    df_val.reset_index(drop=True, inplace=True)
+    
+    index_train = df_train.index.tolist()
+    index_val = df_val.index.tolist()
+
+    # --- Use the new VideoDataset classes ---
+    base_data_path = 'GenVideo'
+    
+    train_dataset = VideoDataset(
+        df=df_train,
+        index_list=index_train,
+        base_data_path=base_data_path,
+        transform=get_train_transforms(),
+        select_frame_nums=8
+    )
+    
+    val_dataset = VideoDatasetVal(
+        df=df_val,
+        index_list=index_val,
+        base_data_path=base_data_path,
+        transform=get_val_transforms(cfg),
+        select_frame_nums=8
+    )
+
+    train_loader = DataLoader(
+        train_dataset, batch_size=cfg['train_batch_size'], shuffle=True, 
+        num_workers=cfg['num_workers'], pin_memory=True, drop_last=True
+    )
+
+    val_loader = DataLoader(
+        val_dataset, batch_size=cfg['val_batch_size'], shuffle=False, 
+        num_workers=cfg['num_workers'], pin_memory=True, drop_last=False
+    )
+    
+    print(f"******* Training Videos {len(index_train)}, Batch size {cfg['train_batch_size']} *******")
+    print(f"******* Testing Videos {len(index_val)}, Batch size {cfg['val_batch_size']} *******")
+
+    return train_loader, val_loader

eval.py

@@ -0,0 +1,48 @@
+import models
+import time
+import torch
+import math
+from tqdm import tqdm
+import numpy as np
+import pandas as pd
+from sklearn.metrics import f1_score, accuracy_score, recall_score, precision_score, average_precision_score, roc_auc_score
+
+def eval_model(cfg, model, val_loader, loss_ce, val_batch_size):
+    model.eval()
+    outpred_list = []
+    gt_label_list = []
+    video_list = []
+    valLoss = 0
+    lossTrainNorm = 0
+    print("******** Start Testing. ********")
+
+    with torch.no_grad():  # No need to track gradients during validation
+        for i, (_, input, target, binary_label, video_id) in enumerate(tqdm(val_loader, desc="Validation", total=len(val_loader))):
+            if i == 0:
+                ss_time = time.time()
+            
+            input = input[:,0]
+            varInput = torch.autograd.Variable(input.float().cuda())
+            varTarget = torch.autograd.Variable(target.contiguous().cuda())
+            var_Binary_Target = torch.autograd.Variable(binary_label.contiguous().cuda())
+
+            logit = model(varInput)
+            lossvalue = loss_ce(logit, var_Binary_Target)
+
+            valLoss += lossvalue.item()
+            lossTrainNorm += 1
+            outpred_list.append(logit[:,0].sigmoid().cpu().detach().numpy())
+            gt_label_list.append(varTarget.cpu().detach().numpy())
+            video_list.append(video_id)
+    
+    valLoss = valLoss / lossTrainNorm
+
+    outpred = np.concatenate(outpred_list, 0)
+    gt_label = np.concatenate(gt_label_list, 0)
+    video_list = np.concatenate(video_list, 0)
+    pred_labels = [1 if item > 0.5 else 0 for item in outpred]
+    true_labels = np.argmax(gt_label, axis=1)
+
+    pred_accuracy = accuracy_score(true_labels, pred_labels)
+
+    return pred_accuracy, video_list, pred_labels, true_labels, outpred

eval2.py

@@ -0,0 +1,110 @@
+import os
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from torchvision import transforms
+from PIL import Image
+import torch.nn.functional as F
+from sklearn.metrics import accuracy_score, average_precision_score, roc_auc_score
+
+import models  # assuming models.py is in your PYTHONPATH or same dir
+
+# -------- Build model as per your code ----------
+def build_model(model_name):
+    if model_name == 'F3Net':
+        model = models.Det_F3_Net()
+    elif model_name == 'NPR':
+        model = models.resnet50_npr()
+    elif model_name == 'STIL':
+        model = models.Det_STIL()
+    elif model_name == 'XCLIP_DeMamba':
+        model = models.XCLIP_DeMamba()
+    elif model_name == 'CLIP_DeMamba':
+        model = models.CLIP_DeMamba()
+    elif model_name == 'XCLIP':
+        model = models.XCLIP()
+    elif model_name == 'CLIP':
+        model = models.CLIP_Base()
+    elif model_name == 'ViT_B_MINTIME':
+        model = models.ViT_B_MINTIME()
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+    return model
+
+
+# --------- Evaluation loop -------------
+def eval_on_frames(model, frames_dir, device):
+    model.eval()
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225])
+    ])
+
+    frame_paths = []
+    for root, _, files in os.walk(frames_dir):
+        for f in files:
+            if f.lower().endswith(('.png', '.jpg', '.jpeg')):
+                frame_paths.append(os.path.join(root, f))
+    frame_paths.sort()
+
+    results = []
+
+    with torch.no_grad():
+        for fp in tqdm(frame_paths, desc="Evaluating frames"):
+            img = Image.open(fp).convert("RGB")
+            
+            # Transform and add batch dimension
+            x = transform(img).unsqueeze(0).to(device)  # [1, C, H, W]
+            
+            # Add temporal dimension expected by DeMamba/XCLIP (T=8)
+            x = x.unsqueeze(1).repeat(1, 8, 1, 1, 1)    # [1, 8, C, H, W]
+
+            # Forward pass
+            logit = model(x)
+            prob = torch.sigmoid(logit[:, 0]).item()
+            pred_label = int(prob > 0.5)
+
+            results.append({
+                "file_name": os.path.basename(fp),
+                "predicted_prob": prob,
+                "predicted_label": pred_label
+            })
+
+    return pd.DataFrame(results)
+
+
+if __name__ == "__main__":
+    # ----- config -----
+    model_name = "XCLIP_DeMamba"  # change if needed
+    model_path = "/home/kalpit/workspace/aigc/repos/SAFE_challenge/seetrails_aigvdet_v2.0.0/results/kling_9k_9k/best_acc.pth"
+    frames_dir = "./frames"
+    output_csv = "results.csv"
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # ---- load model ----
+    model = build_model(model_name).to(device)
+    checkpoint = torch.load(model_path, map_location=device)
+    model.load_state_dict(checkpoint["model_state_dict"], strict=False)
+    print(f"Loaded model weights from {model_path}")
+
+    # ---- evaluate ----
+    df_results = eval_on_frames(model, frames_dir, device)
+    df_results.to_csv(output_csv, index=False)
+    print(f"Saved framewise results to {output_csv}")
+
+    # ---- optional: basic metrics if you have GT ----
+    if "label" in df_results.columns:
+        y_true = df_results["label"].values
+        y_pred = df_results["predicted_label"].values
+        y_prob = df_results["predicted_prob"].values
+
+        acc = accuracy_score(y_true, y_pred)
+        auc = roc_auc_score(y_true, y_prob)
+        ap = average_precision_score(y_true, y_prob)
+
+        print(f"Accuracy: {acc:.4f}, AUC: {auc:.4f}, AP: {ap:.4f}")

extract_frames.py

@@ -0,0 +1,74 @@
+import os
+import pandas as pd
+from datasets import load_dataset
+import numpy as np
+import tqdm.auto as tqdm
+import io
+import torch
+import av
+from torchvision.utils import save_image
+
+
+def preprocess_and_save_frames(
+    file_like: io.BytesIO,
+    video_id: str,
+    output_root: str = "./frames",
+    crop_size: int = -1,
+    every: int = 10,
+    max_memory: int = 50 * 1024 * 1024,
+    device: str = "cpu"
+):
+    """
+    Loads a video and saves frames as images in output_root/<video_id>/.
+    """
+
+    # Ensure the base frames directory exists
+    os.makedirs(output_root, exist_ok=True)
+
+    # Create subfolder for this specific video
+    video_dir = os.path.join(output_root, video_id)
+    os.makedirs(video_dir, exist_ok=True)
+
+    center_crop_transform = None
+    if crop_size > 0:
+        from torchvision import transforms
+        center_crop_transform = transforms.CenterCrop(crop_size)
+
+    file_like.seek(0)
+    container = av.open(file_like)
+    current_memory = 0
+
+    for i, frame in enumerate(container.decode(video=0)):
+        if i % every == 0:
+            frame_array = frame.to_ndarray(format="rgb24")
+            frame_tensor = torch.from_numpy(frame_array).permute(2, 0, 1).float() / 255.0
+
+            if center_crop_transform is not None:
+                frame_tensor = center_crop_transform(frame_tensor)
+
+            frame_path = os.path.join(video_dir, f"frame_{i:05d}.png")
+            save_image(frame_tensor, frame_path)
+
+            frame_bytes = frame_tensor.numel() * 4
+            current_memory += frame_bytes
+            if current_memory >= max_memory:
+                break
+
+
+# -------- Main section --------
+if __name__ == "__main__":
+    DATASET_PATH = "/tmp/data"
+    dataset_remote = load_dataset(DATASET_PATH, split="test", streaming=True)
+
+    # Make sure ./frames exists before processing
+    os.makedirs("./frames", exist_ok=True)
+
+    for el in tqdm.tqdm(dataset_remote, desc="Extracting frames"):
+        try:
+            video_id = str(el["id"])
+            file_like = io.BytesIO(el["video"]["bytes"])
+            preprocess_and_save_frames(file_like, video_id)
+        except Exception as e:
+            print(f"⚠️ Failed {el['id']}: {e}")
+
+    print("✅ All frames saved under ./frames/<video_id>/")

models/DeMamba.py

@@ -0,0 +1,176 @@
+from transformers import XCLIPVisionModel
+import os
+import sys
+import numpy as np
+import torch
+import torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+from .mamba_base import MambaConfig, ResidualBlock
+import torch.nn.init as init
+from .clip import clip
+import math
+from transformers import XCLIPVisionConfig, XCLIPVisionModel
+
+def create_reorder_index(N, device):
+    new_order = []
+    for col in range(N):
+        if col % 2 == 0:
+            new_order.extend(range(col, N*N, N))
+        else:
+            new_order.extend(range(col + N*(N-1), col-1, -N))
+    return torch.tensor(new_order, device=device)
+
+def reorder_data(data, N):
+    assert isinstance(data, torch.Tensor), "data should be a torch.Tensor"
+    device = data.device
+    new_order = create_reorder_index(N, device)
+    B, t, _, _ = data.shape
+    index = new_order.repeat(B, t, 1).unsqueeze(-1)
+    reordered_data = torch.gather(data, 2, index.expand_as(data))
+    return reordered_data
+
+class XCLIP_DeMamba(nn.Module):
+    def __init__(
+        self, channel_size=768, class_num=1
+    ):
+        super(XCLIP_DeMamba, self).__init__()
+        # self.encoder = XCLIPVisionModel.from_pretrained("GenVideo/pretrained_weights/xclip")
+        # my code for training from scratch
+        config = XCLIPVisionConfig()
+        self.encoder = XCLIPVisionModel(config)
+        
+        blocks = []
+        channel = 768
+        self.fusing_ratios = 1
+        self.patch_nums = (14//self.fusing_ratios)**2
+        self.mamba_configs = MambaConfig(d_model=channel)
+        self.mamba = ResidualBlock(config = self.mamba_configs)
+        # self.fc1 = nn.Linear((self.patch_nums+1)*channel, class_num)
+        self.fc1 = nn.Linear(38400, class_num) # my code
+        # self.fc_norm = nn.LayerNorm(self.patch_nums*channel)
+        self.fc_norm = None # my code
+        self.fc_norm2 = nn.LayerNorm(768)
+        self.initialize_weights(self.fc1)
+        self.dropout = nn.Dropout(p=0.0)
+
+    def initialize_weights(self, module):
+        for m in module.modules():
+            if isinstance(m, nn.Linear):
+                init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Conv2d):
+                init.kaiming_uniform_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        b, t, _, h, w = x.shape
+        images = x.view(b * t, 3, h, w)
+        outputs = self.encoder(images, output_hidden_states=True)
+        sequence_output = outputs['last_hidden_state'][:,1:,:]
+        _, _, c = sequence_output.shape
+
+        global_feat = outputs['pooler_output'].reshape(b, t, -1)
+        global_feat = global_feat.mean(1)
+        global_feat = self.fc_norm2(global_feat)
+
+        sequence_output = sequence_output.view(b, t, -1, c)
+        _, _, f_w, _ = sequence_output.shape
+        f_h, f_w = int(math.sqrt(f_w)), int(math.sqrt(f_w))
+
+        s = f_h//self.fusing_ratios
+        sequence_output = sequence_output.view(b, t, self.fusing_ratios, s, self.fusing_ratios, s, c)
+        x = sequence_output.permute(0, 2, 4, 1, 3, 5, 6).contiguous().view(b*s*s, t, -1, c)
+        b_l = b*s*s
+        
+        x = reorder_data(x, self.fusing_ratios)
+        x = x.permute(0, 2, 1, 3).contiguous().view(b_l, -1, c)
+        res = self.mamba(x)
+
+        video_level_features = res.mean(1)
+        video_level_features = video_level_features.view(b, -1)
+        
+        # my code
+        if self.fc_norm is None:
+            self.fc_norm = nn.LayerNorm(video_level_features.size(-1)).to(video_level_features.device)
+        
+        video_level_features = self.fc_norm(video_level_features)
+        video_level_features = torch.cat((global_feat, video_level_features), dim=1)
+
+        pred = self.fc1(video_level_features)
+        pred = self.dropout(pred)
+
+        return pred
+
+
+
+class CLIP_DeMamba(nn.Module):
+    def __init__(
+        self, channel_size=512, class_num=1
+    ):
+        super(CLIP_DeMamba, self).__init__()
+        self.clip_model, preprocess = clip.load('ViT-B-14')
+        self.clip_model = self.clip_model.float()
+        blocks = []
+        channel = 512
+        self.fusing_ratios = 2
+        self.patch_nums = (14//self.fusing_ratios)**2
+        self.mamba_configs = MambaConfig(d_model=channel)
+        self.mamba = ResidualBlock(config = self.mamba_configs)
+        
+        self.fc1 = nn.Linear(channel*(self.patch_nums+1), class_num)
+        
+        self.bn1 = nn.BatchNorm1d(channel)
+        self.initialize_weights(self.fc1)
+
+    def initialize_weights(self, module):
+        for m in module.modules():
+            if isinstance(m, nn.Linear):
+                init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Conv2d):
+                init.kaiming_uniform_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        b, t, _, h, w = x.shape
+        images = x.view(b * t, 3, h, w)
+        sequence_output = self.clip_model.encode_image(images)
+        _, _, c = sequence_output.shape
+        sequence_output = sequence_output.view(b, t, -1, c)
+
+        global_feat = sequence_output.reshape(b, -1, c)
+        global_feat = global_feat.mean(1)
+
+        _, _, f_w, _ = sequence_output.shape
+        f_h, f_w = int(math.sqrt(f_w)), int(math.sqrt(f_w))
+
+        s = f_h//self.fusing_ratios
+        sequence_output = sequence_output.view(b, t, self.fusing_ratios, s, self.fusing_ratios, s, c)
+        x = sequence_output.permute(0, 2, 4, 1, 3, 5, 6).contiguous().view(b*s*s, t, -1, c)
+        b_l = b*s*s
+        
+        x = reorder_data(x, self.fusing_ratios)
+        x = x.permute(0, 2, 1, 3).contiguous().view(b_l, -1, c)
+        res = self.mamba(x)
+        video_level_features = res.mean(1)
+        video_level_features = video_level_features.view(b, -1)
+
+        video_level_features = torch.cat((global_feat, video_level_features), dim=1)
+        x = self.fc1(video_level_features)
+
+        return x
+
+if __name__ == '__main__':
+    model = CLIP_DeMamba()
+    print(model)

models/F3Net.py

@@ -0,0 +1,434 @@
+"""
+F3Net: Fusion, Feedback and Focus for Salient Object Detection @ AAAI'2020
+Copyright (c) University of Chinese Academy of Sciences and its affiliates.
+Modified by Jun Wei from https://github.com/weijun88/F3Net
+"""
+import os
+import sys
+import numpy as np
+import torch
+import torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+
+pretrained_settings = {
+    'xception': {
+        'imagenet': {
+            'dir': '/ossfs/workspace/aigc_video/weights/xception-b5690688.pth',
+            'url': 'http://data.lip6.fr/cadene/pretrainedmodels/xception-b5690688.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 299, 299],
+            'input_range': [0, 1],
+            'mean': [0.5, 0.5, 0.5],
+            'std': [0.5, 0.5, 0.5],
+            'num_classes': 1000,
+            'scale': 0.8975 # The resize parameter of the validation transform should be 333, and make sure to center crop at 299x299
+        }
+    }
+}
+
+class F3Net(nn.Module):
+    """
+    Implementation is mainly referenced from https://github.com/yyk-wew/F3Net
+    """
+    def __init__(self, 
+                 num_classes: int=2, 
+                 img_width: int=299, 
+                 img_height: int=299, 
+                 LFS_window_size: int=10, 
+                 LFS_M: int=6) -> None:
+        super(F3Net, self).__init__()
+        assert img_width == img_height
+        self.img_size = img_width
+        self.num_classes = num_classes
+        self._LFS_window_size = LFS_window_size
+        self._LFS_M = LFS_M
+        
+        
+        self.fad_head = FAD_Head(self.img_size)
+        self.lfs_head = LFS_Head(self.img_size, self._LFS_window_size, self._LFS_M)
+        
+        self.fad_excep = self._init_xcep_fad()
+        self.lfs_excep = self._init_xcep_lfs()
+        
+        self.mix_block7 = MixBlock(c_in=728, width=19, height=19) 
+        self.mix_block12 = MixBlock(c_in=1024, width=10, height=10) 
+        self.excep_forwards = ['conv1', 'bn1', 'relu', 'conv2', 'bn2', 'relu', 
+                               'block1', 'block2', 'block3', 'block4', 'block5', 'block6', 
+                               'block7', 'block8', 'block9', 'block10' , 'block11', 'block12',
+                               'conv3', 'bn3', 'relu', 'conv4', 'bn4']
+
+         # classifier
+        self.relu = nn.ReLU(inplace=True)
+        self.fc = nn.Linear(4096, num_classes)
+        self.dp = nn.Dropout(p=0.2)
+        
+    def _init_xcep_fad(self):
+        fad_excep =  return_pytorch04_xception(True)
+        conv1_data = fad_excep.conv1.weight.data
+        # let new conv1 use old param to balance the network
+        fad_excep.conv1 = nn.Conv2d(12, 32, 3, 2, 0, bias=False)
+        for i in range(4):
+            fad_excep.conv1.weight.data[:, i*3:(i+1)*3, :, :] = conv1_data / 4.0
+        return fad_excep
+    
+    def  _init_xcep_lfs(self): 
+        lfs_excep = return_pytorch04_xception(True)
+        conv1_data = lfs_excep.conv1.weight.data
+        # let new conv1 use old param to balance the network
+        lfs_excep.conv1 = nn.Conv2d(self._LFS_M, 32, 3, 1, 0, bias=False)
+        for i in range(int(self._LFS_M / 3)):
+            lfs_excep.conv1.weight.data[:, i*3:(i+1)*3, :, :] = conv1_data / float(self._LFS_M / 3.0)
+        return lfs_excep
+    
+    def _features(self, x_fad, x_fls):
+        for forward_func in self.excep_forwards:
+            x_fad = getattr(self.fad_excep, forward_func)(x_fad)
+            x_fls = getattr(self.lfs_excep, forward_func)(x_fls)
+            if forward_func == 'block7':
+                x_fad, x_fls = self.mix_block7(x_fad, x_fls)
+            if forward_func == 'block12':
+                x_fad, x_fls = self.mix_block12(x_fad, x_fls)
+        return x_fad, x_fls
+    
+    def _norm_feature(self, x):
+        x = self.relu(x)
+        x = F.adaptive_avg_pool2d(x, (1,1))
+        x = x.view(x.size(0), -1)
+        return x
+    
+    def forward(self, x):
+        fad_input = self.fad_head(x)
+        lfs_input = self.lfs_head(x)
+        x_fad, x_fls = self._features(fad_input, lfs_input)
+        x_fad = self._norm_feature(x_fad)
+        x_fls = self._norm_feature(x_fls)
+        x_cat = torch.cat((x_fad, x_fls), dim=1)
+        x_drop = self.dp(x_cat)
+        logit = self.fc(x_drop)
+        return logit
+
+class SeparableConv2d(nn.Module):
+    def __init__(self,in_channels,out_channels,kernel_size=1,stride=1,padding=0,dilation=1,bias=False):
+        super(SeparableConv2d,self).__init__()
+
+        self.conv1 = nn.Conv2d(in_channels,in_channels,kernel_size,stride,padding,dilation,groups=in_channels,bias=bias)
+        self.pointwise = nn.Conv2d(in_channels,out_channels,1,1,0,1,1,bias=bias)
+
+    def forward(self,x):
+        x = self.conv1(x)
+        x = self.pointwise(x)
+        return x
+
+class Block(nn.Module):
+    def __init__(self,in_filters,out_filters,reps,strides=1,start_with_relu=True,grow_first=True):
+        super(Block, self).__init__()
+
+        if out_filters != in_filters or strides!=1:
+            self.skip = nn.Conv2d(in_filters,out_filters,1,stride=strides, bias=False)
+            self.skipbn = nn.BatchNorm2d(out_filters)
+        else:
+            self.skip=None
+
+        self.relu = nn.ReLU(inplace=True)
+        rep=[]
+
+        filters=in_filters
+        if grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False))
+            rep.append(nn.BatchNorm2d(out_filters))
+            filters = out_filters
+
+        for i in range(reps-1):
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(filters,filters,3,stride=1,padding=1,bias=False))
+            rep.append(nn.BatchNorm2d(filters))
+
+        if not grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False))
+            rep.append(nn.BatchNorm2d(out_filters))
+
+        if not start_with_relu:
+            rep = rep[1:]
+        else:
+            rep[0] = nn.ReLU(inplace=False)
+
+        if strides != 1:
+            rep.append(nn.MaxPool2d(3,strides,1))
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self,inp):
+        x = self.rep(inp)
+
+        if self.skip is not None:
+            skip = self.skip(inp)
+            skip = self.skipbn(skip)
+        else:
+            skip = inp
+
+        x+=skip
+        return x
+
+class Xception(nn.Module):
+    """
+    Xception optimized for the ImageNet dataset, as specified in
+    https://arxiv.org/pdf/1610.02357.pdf
+    """
+    def __init__(self, num_classes=1000):
+        """ Constructor
+        Args:
+            num_classes: number of classes
+        """
+        super(Xception, self).__init__()
+        self.num_classes = num_classes
+
+        self.conv1 = nn.Conv2d(3, 32, 3,2, 0, bias=False)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(32,64,3,bias=False)
+        self.bn2 = nn.BatchNorm2d(64)
+        #do relu here
+
+        self.block1=Block(64,128,2,2,start_with_relu=False,grow_first=True)
+        self.block2=Block(128,256,2,2,start_with_relu=True,grow_first=True)
+        self.block3=Block(256,728,2,2,start_with_relu=True,grow_first=True)
+
+        self.block4=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+        self.block5=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+        self.block6=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+        self.block7=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+
+        self.block8=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+        self.block9=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+        self.block10=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+        self.block11=Block(728,728,3,1,start_with_relu=True,grow_first=True)
+
+        self.block12=Block(728,1024,2,2,start_with_relu=True,grow_first=False)
+
+        self.conv3 = SeparableConv2d(1024,1536,3,1,1)
+        self.bn3 = nn.BatchNorm2d(1536)
+
+        #do relu here
+        self.conv4 = SeparableConv2d(1536,2048,3,1,1)
+        self.bn4 = nn.BatchNorm2d(2048)
+
+def xception(num_classes=1000, pretrained='imagenet'):
+    model = Xception(num_classes=num_classes)
+    if pretrained:
+        settings = pretrained_settings['xception'][pretrained]
+        assert num_classes == settings['num_classes'], \
+            "num_classes should be {}, but is {}".format(settings['num_classes'], num_classes)
+
+        model = Xception(num_classes=num_classes)
+        model.load_state_dict(settings['dir'])
+
+        model.input_space = settings['input_space']
+        model.input_size = settings['input_size']
+        model.input_range = settings['input_range']
+        model.mean = settings['mean']
+        model.std = settings['std']
+
+    return model
+
+def return_pytorch04_xception(pretrained=True):
+    model = xception(pretrained=False)
+    if pretrained:
+        state_dict = torch.load(
+            '/ossfs/workspace/GenVideo/weights/xception-b5690688.pth')
+        for name, weights in state_dict.items():
+            if 'pointwise' in name:
+                state_dict[name] = weights.unsqueeze(-1).unsqueeze(-1)
+        model.load_state_dict(state_dict, strict=False)
+ 
+    return model
+
+class Filter(nn.Module):
+    def __init__(self, size, 
+                 band_start, 
+                 band_end, 
+                 use_learnable=True, 
+                 norm=False):
+        super(Filter, self).__init__()
+        self.use_learnable = use_learnable
+
+        self.base = nn.Parameter(torch.tensor(generate_filter(band_start, band_end, size)), requires_grad=False)
+        if self.use_learnable:
+            self.learnable = nn.Parameter(torch.randn(size, size), requires_grad=True)
+            self.learnable.data.normal_(0., 0.1)
+            # Todo
+            # self.learnable = nn.Parameter(torch.rand((size, size)) * 0.2 - 0.1, requires_grad=True)
+
+        self.norm = norm
+        if norm:
+            self.ft_num = nn.Parameter(torch.sum(torch.tensor(generate_filter(band_start, band_end, size))), requires_grad=False)
+
+
+    def forward(self, x):
+        if self.use_learnable:
+            filt = self.base + norm_sigma(self.learnable)
+        else:
+            filt = self.base
+
+        if self.norm:
+            y = x * filt / self.ft_num
+        else:
+            y = x * filt
+        return y
+
+class FAD_Head(nn.Module):
+    def __init__(self, size):
+        super(FAD_Head, self).__init__()
+        # init DCT matrix
+        self._DCT_all = nn.Parameter(torch.tensor(DCT_mat(size)).float(), requires_grad=False)
+        self._DCT_all_T = nn.Parameter(torch.transpose(torch.tensor(DCT_mat(size)).float(), 0, 1), requires_grad=False)
+
+        # define base filters and learnable
+        # 0 - 1/16 || 1/16 - 1/8 || 1/8 - 1
+        low_filter = Filter(size, 0, size // 16)
+        middle_filter = Filter(size, size // 16, size // 8)
+        high_filter = Filter(size, size // 8, size)
+        all_filter = Filter(size, 0, size * 2)
+
+        self.filters = nn.ModuleList([low_filter, middle_filter, high_filter, all_filter])
+
+    def forward(self, x):
+        # DCT
+        x_freq = self._DCT_all @ x @ self._DCT_all_T    # [N, 3, 299, 299]
+
+        # 4 kernel
+        y_list = []
+        for i in range(4):
+            x_pass = self.filters[i](x_freq)  # [N, 3, 299, 299]
+            y = self._DCT_all_T @ x_pass @ self._DCT_all    # [N, 3, 299, 299]
+            y_list.append(y)
+        out = torch.cat(y_list, dim=1)    # [N, 12, 299, 299]
+        return out
+
+
+class LFS_Head(nn.Module):
+    def __init__(self, size, window_size, M):
+        super(LFS_Head, self).__init__()
+
+        self.window_size = window_size
+        self._M = M
+
+        # init DCT matrix
+        self._DCT_patch = nn.Parameter(torch.tensor(DCT_mat(window_size)).float(), requires_grad=False)
+        self._DCT_patch_T = nn.Parameter(torch.transpose(torch.tensor(DCT_mat(window_size)).float(), 0, 1), requires_grad=False)
+
+        self.unfold = nn.Unfold(kernel_size=(window_size, window_size), stride=2, padding=4)
+
+        # init filters
+        self.filters = nn.ModuleList([Filter(window_size, window_size * 2. / M * i, window_size * 2. / M * (i+1), norm=True) for i in range(M)])
+    
+    def forward(self, x):
+        # turn RGB into Gray
+        x_gray = 0.299*x[:,0,:,:] + 0.587*x[:,1,:,:] + 0.114*x[:,2,:,:]
+        x = x_gray.unsqueeze(1)
+
+        # rescale to 0 - 255
+        x = (x + 1.) * 122.5
+
+        # calculate size
+        N, C, W, H = x.size()
+        S = self.window_size
+        size_after = int((W - S + 8)/2) + 1
+        assert size_after == 149
+
+        # sliding window unfold and DCT
+        x_unfold = self.unfold(x)   # [N, C * S * S, L]   L:block num
+        L = x_unfold.size()[2]
+        x_unfold = x_unfold.transpose(1, 2).reshape(N, L, C, S, S)  # [N, L, C, S, S]
+        x_dct = self._DCT_patch @ x_unfold @ self._DCT_patch_T
+
+        # M kernels filtering
+        y_list = []
+        for i in range(self._M):
+            # y = self.filters[i](x_dct)    # [N, L, C, S, S]
+            # y = torch.abs(y)
+            # y = torch.sum(y, dim=[2,3,4])   # [N, L]
+            # y = torch.log10(y + 1e-15)
+            y = torch.abs(x_dct)
+            y = torch.log10(y + 1e-15)
+            y = self.filters[i](y)
+            y = torch.sum(y, dim=[2,3,4])
+            y = y.reshape(N, size_after, size_after).unsqueeze(dim=1)   # [N, 1, 149, 149]
+            y_list.append(y)
+        out = torch.cat(y_list, dim=1)  # [N, M, 149, 149]
+        return out
+    
+class MixBlock(nn.Module):
+    
+    def __init__(self, c_in, width, height):
+        super(MixBlock, self).__init__()
+        self.FAD_query = nn.Conv2d(c_in, c_in, (1,1))
+        self.LFS_query = nn.Conv2d(c_in, c_in, (1,1))
+
+        self.FAD_key = nn.Conv2d(c_in, c_in, (1,1))
+        self.LFS_key = nn.Conv2d(c_in, c_in, (1,1))
+
+        self.softmax = nn.Softmax(dim=-1)
+        self.relu = nn.ReLU()
+
+        self.FAD_gamma = nn.Parameter(torch.zeros(1))
+        self.LFS_gamma = nn.Parameter(torch.zeros(1))
+
+        self.FAD_conv = nn.Conv2d(c_in, c_in, (1,1), groups=c_in)
+        self.FAD_bn = nn.BatchNorm2d(c_in)
+        self.LFS_conv = nn.Conv2d(c_in, c_in, (1,1), groups=c_in)
+        self.LFS_bn = nn.BatchNorm2d(c_in)
+
+    def forward(self, x_FAD, x_LFS):
+        B, C, W, H = x_FAD.size()
+        assert W == H
+
+        q_FAD = self.FAD_query(x_FAD).view(-1, W, H)    # [BC, W, H]
+        q_LFS = self.LFS_query(x_LFS).view(-1, W, H)
+        M_query = torch.cat([q_FAD, q_LFS], dim=2)  # [BC, W, 2H]
+
+        k_FAD = self.FAD_key(x_FAD).view(-1, W, H).transpose(1, 2)  # [BC, H, W]
+        k_LFS = self.LFS_key(x_LFS).view(-1, W, H).transpose(1, 2)
+        M_key = torch.cat([k_FAD, k_LFS], dim=1)    # [BC, 2H, W]
+
+        energy = torch.bmm(M_query, M_key)  #[BC, W, W]
+        attention = self.softmax(energy).view(B, C, W, W)
+
+        att_LFS = x_LFS * attention * (torch.sigmoid(self.LFS_gamma) * 2.0 - 1.0)
+        y_FAD = x_FAD + self.FAD_bn(self.FAD_conv(att_LFS))
+
+        att_FAD = x_FAD * attention * (torch.sigmoid(self.FAD_gamma) * 2.0 - 1.0)
+        y_LFS = x_LFS + self.LFS_bn(self.LFS_conv(att_FAD))
+        return y_FAD, y_LFS
+
+def DCT_mat(size):
+    m = [[ (np.sqrt(1./size) if i == 0 else np.sqrt(2./size)) * np.cos((j + 0.5) * np.pi * i / size) for j in range(size)] for i in range(size)]
+    return m
+
+def generate_filter(start, end, size):
+    return [[0. if i + j > end or i + j <= start else 1. for j in range(size)] for i in range(size)]
+
+def norm_sigma(x):
+    return 2. * torch.sigmoid(x) - 1.
+
+
+class Det_F3_Net(nn.Module):
+    def __init__(self):
+        super(Det_F3_Net, self).__init__()
+        self.f3net = F3Net(num_classes=1)
+
+    def forward(self, x):
+        b, t, _, h, w = x.shape
+        images = x.view(b * t, 3, h, w)
+        sequence_output = self.f3net(images)
+        sequence_output = sequence_output.view(b, t, -1)
+        sequence_output = sequence_output.mean(1)
+
+        return sequence_output
+
+if __name__ == '__main__':
+    model = F3Net()
+    print(model)

models/FTCN.py

@@ -0,0 +1,143 @@
+'''
+Exploring Temporal Coherence for More General Video Face Forgery Detection @ ICCV'2021
+Copyright (c) Xiamen University and its affiliates.
+Modified by Yinglin Zheng from https://github.com/yinglinzheng/FTCN
+'''
+
+import torch
+from torch import nn
+from .time_transformer import TimeTransformer
+from .clip import clip
+
+class RandomPatchPool(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        # batch,channel,16,7x7
+        b, c, t, h, w = x.shape
+        x = x.reshape(b, c, t, h * w)
+        if self.training and my_cfg.model.transformer.random_select:
+            while True:
+                idx = random.randint(0, h * w - 1)
+                i = idx // h
+                j = idx % h
+                if j == 0 or i == h - 1 or j == h - 1:
+                    continue
+                else:
+                    break
+        else:
+            idx = h * w // 2
+        x = x[..., idx]
+        return x
+
+
+def valid_idx(idx, h):
+    i = idx // h
+    j = idx % h
+    if j == 0 or i == h - 1 or j == h - 1:
+        return False
+    else:
+        return True
+
+
+class RandomAvgPool(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        # batch,channel,16,7x7
+        b, c, t, h, w = x.shape
+        x = x.reshape(b, c, t, h * w)
+        candidates = list(range(h * w))
+        candidates = [idx for idx in candidates if valid_idx(idx, h)]
+        max_k = len(candidates)
+        if self.training and my_cfg.model.transformer.random_select:
+            k = my_cfg.model.transformer.k
+        else:
+            k = max_k
+        candidates = random.sample(candidates, k)
+        x = x[..., candidates].mean(-1)
+        return x
+
+class TransformerHead(nn.Module):
+    def __init__(self, spatial_size=7, time_size=8, in_channels=2048):
+        super().__init__()
+        # if my_cfg.model.inco.no_time_pool:
+        #     time_size = time_size * 2
+        patch_type = 'time'
+        if patch_type == "time":
+            self.pool = nn.AvgPool3d((1, spatial_size, spatial_size))
+            self.num_patches = time_size
+        elif patch_type == "spatial":
+            self.pool = nn.AvgPool3d((time_size, 1, 1))
+            self.num_patches = spatial_size ** 2
+        elif patch_type == "random":
+            self.pool = RandomPatchPool()
+            self.num_patches = time_size
+        elif patch_type == "random_avg":
+            self.pool = RandomAvgPool()
+            self.num_patches = time_size
+        elif patch_type == "all":
+            self.pool = nn.Identity()
+            self.num_patches = time_size * spatial_size * spatial_size
+        else:
+            raise NotImplementedError(patch_type)
+
+        self.dim = -1
+        if self.dim == -1:
+            self.dim = in_channels
+
+        self.in_channels = in_channels
+
+        if self.dim != self.in_channels:
+            self.fc = nn.Linear(self.in_channels, self.dim)
+
+        default_params = dict(
+            dim=self.dim, depth=6, heads=16, mlp_dim=2048, dropout=0.1, emb_dropout=0.1,
+        )
+        
+        self.time_T = TimeTransformer(
+            num_patches=self.num_patches, num_classes=1, **default_params
+        )
+
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = x.reshape(-1, self.in_channels, self.num_patches)
+        x = x.permute(0, 2, 1)
+        if self.dim != self.in_channels:
+            x = self.fc(x.reshape(-1, self.in_channels))
+            x = x.reshape(-1, self.num_patches, self.dim)
+        x = self.time_T(x)
+
+        return x
+
+
+class ViT_B_FTCN(nn.Module):
+    def __init__(
+        self, channel_size=512, class_num=1
+    ):
+        super(ViT_B_FTCN, self).__init__()
+        self.clip_model, preprocess = clip.load('ViT-B-16')
+        self.clip_model = self.clip_model.float()
+
+        self.head = TransformerHead(spatial_size=14, time_size=8, in_channels=512)
+
+    def forward(self, x):
+        b, t, _, h, w = x.shape
+        images = x.view(b * t, 3, h, w)
+        sequence_output = self.clip_model.encode_image(images)
+        _, _, c = sequence_output.shape
+        sequence_output = sequence_output.view(b, t, 14, 14, c)
+        sequence_output = sequence_output.permute(0, 4, 1, 2, 3)
+
+        res = self.head(sequence_output)
+
+        return res
+if __name__ == '__main__':
+    model = ViT_B_FTCN()
+    model = model.cuda()
+    dummy_input = torch.randn(4,8,3,224,224)
+    dummy_input = dummy_input.cuda()
+    model(dummy_input)

models/MINTIME

@@ -0,0 +1,269 @@
+"""
+MINTIME: Multi-Identity size-iNvariant TIMEsformer for Video Deepfake Detection@TIFS'2024
+Copyright (c) ISTI-CNR and its affiliates.
+Modified by Davide Alessandro Coccomini from https://github.com/davide-coccomini/MINTIME-Multi-Identity-size-iNvariant-TIMEsformer-for-Video-Deepfake-Detection
+"""
+
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from statistics import mean
+from torch.nn.init import trunc_normal_
+import cv2
+import numpy as np
+from random import random
+from .clip import clip
+from einops.layers.torch import Rearrange
+
+# helpers
+def exists(val):
+    return val is not None
+
+# classes
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, *args, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, *args, **kwargs)
+
+# time token shift
+def shift(t, amt):
+    if amt == 0:
+        return t
+    return F.pad(t, (0, 0, 0, 0, amt, -amt))
+
+class PreTokenShift(nn.Module):
+    def __init__(self, frames, fn):
+        super().__init__()
+        self.frames = frames
+        self.fn = fn
+
+    def forward(self, x, *args, **kwargs):
+        f, dim = self.frames, x.shape[-1]
+        cls_x, x = x[:, :1], x[:, 1:]
+        x = rearrange(x, 'b (f n) d -> b f n d', f = f)
+
+        # shift along time frame before and after
+        dim_chunk = (dim // 3)
+        chunks = x.split(dim_chunk, dim = -1)
+        chunks_to_shift, rest = chunks[:3], chunks[3:]
+        shifted_chunks = tuple(map(lambda args: shift(*args), zip(chunks_to_shift, (-1, 0, 1))))
+        x = torch.cat((*shifted_chunks, *rest), dim = -1)
+
+        x = rearrange(x, 'b f n d -> b (f n) d')
+        x = torch.cat((cls_x, x), dim = 1)
+        return self.fn(x, *args, **kwargs)
+
+# feedforward
+class GEGLU(nn.Module):
+    def forward(self, x):
+        x, gates = x.chunk(2, dim = -1)
+        return x * F.gelu(gates)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult = 4, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult * 2),
+            GEGLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+# attention
+def attn(q, k, v):
+    sim = torch.einsum('b i d, b j d -> b i j', q, k)
+    attn = sim.softmax(dim = -1)
+    out = torch.einsum('b i j, b j d -> b i d', attn, v)
+    return out, attn
+
+class Attention(nn.Module):
+    def __init__(self, dim, dim_head = 64, heads = 8, dropout = 0.):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        inner_dim = dim_head * heads
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x, einops_from, einops_to, **einops_dims):
+        h = self.heads
+        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))
+
+        q = q * self.scale
+
+        # splice out classification token at index 1
+        (cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, :1], t[:, 1:]), (q, k, v))
+
+        # let classification token attend to key / values of all patches across time and space
+        cls_out, cls_attentions = attn(cls_q, k, v)
+
+        # rearrange across time or space
+        q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), (q_, k_, v_))
+
+        # expand cls token keys and values across time or space and concat
+        r = q_.shape[0] // cls_k.shape[0]
+        cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r = r), (cls_k, cls_v))
+
+        k_ = torch.cat((cls_k, k_), dim = 1)
+        v_ = torch.cat((cls_v, v_), dim = 1)
+
+        # attention
+        out, attentions = attn(q_, k_, v_)
+
+        # merge back time or space
+        out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)
+
+        # concat back the cls token
+        out = torch.cat((cls_out, out), dim = 1)
+
+        # merge back the heads
+        out = rearrange(out, '(b h) n d -> b n (h d)', h = h)
+
+        # combine heads out
+        return self.to_out(out), cls_attentions
+
+class SizeInvariantTimeSformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        require_attention = False
+    ):
+        super().__init__()
+        self.dim = 512
+        self.num_frames = 8
+        self.max_identities = 1
+        self.image_size = 224
+        self.num_classes = 1
+        self.patch_size = 1
+        self.num_patches = 196
+        self.channels = 512
+        self.depth = 9
+        self.heads = 8
+        self.dim_head = 64
+        self.attn_dropout = 0.
+        self.ff_dropout = 0.
+        self.shift_tokens = False
+        self.enable_size_emb = True
+        self.enable_pos_emb = True
+        self.require_attention = require_attention
+
+        num_positions = self.num_frames * self.channels
+        self.to_patch_embedding = nn.Linear(self.channels, self.dim)
+        self.cls_token = nn.Parameter(torch.randn(1, self.dim))
+        self.pos_emb = nn.Embedding(num_positions + 1, self.dim)
+
+        if self.enable_size_emb:
+            self.size_emb = nn.Embedding(num_positions + 1, self.dim)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(self.depth):
+            ff = FeedForward(self.dim, dropout = self.ff_dropout)
+            time_attn = Attention(self.dim, dim_head = self.dim_head, heads = self.heads, dropout = self.attn_dropout)
+            spatial_attn = Attention(self.dim, dim_head = self.dim_head, heads = self.heads, dropout = self.attn_dropout)
+            if self.shift_tokens:
+                time_attn, spatial_attn, ff = map(lambda t: PreTokenShift(self.num_frames, t), (time_attn, spatial_attn, ff))
+
+            time_attn, spatial_attn, ff = map(lambda t: PreNorm(self.dim, t), (time_attn, spatial_attn, ff))
+            self.layers.append(nn.ModuleList([time_attn, spatial_attn, ff]))
+
+        self.to_out = nn.Sequential(
+            nn.LayerNorm(self.dim),
+            nn.Linear(self.dim, self.num_classes)
+        )
+
+        # Initialization
+        trunc_normal_(self.pos_emb.weight, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        if self.enable_size_emb:
+            trunc_normal_(self.size_emb.weight, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_emb', 'cls_token'}
+
+    def forward(self, x):
+        b, f, c, h, w = x.shape
+        n = h * w
+        device = x.device
+
+        x = rearrange(x, 'b f c h w -> b (f h w) c')  # B x F*P*P x C
+        tokens = self.to_patch_embedding(x)  # B x 8*7*7 x dim
+
+        # Add cls token
+        cls_token = repeat(self.cls_token, 'n d -> b n d', b = b)
+        x = torch.cat((cls_token, tokens), dim = 1)
+
+        # Positional embedding
+        x += self.pos_emb(torch.arange(x.shape[1], device=device))
+
+        # Time and space attention
+        for (time_attn, spatial_attn, ff) in self.layers:
+            y, _ = time_attn(x, 'b (f n) d', '(b n) f d', n = n)
+            x = x + y
+            y, _ = spatial_attn(x, 'b (f n) d', '(b f) n d', f = f)
+            x = x + y
+            x = ff(x) + x
+
+        cls_token = x[:, 0]
+
+        if self.require_attention:
+            return self.to_out(cls_token)
+        else:
+            return self.to_out(cls_token)
+
+
+class ViT_B_MINTIME(nn.Module):
+    def __init__(
+        self, channel_size=512, class_num=1
+    ):
+        super(ViT_B_MINTIME, self).__init__()
+        self.clip_model, preprocess = clip.load('ViT-B-16')
+        self.clip_model = self.clip_model.float()
+        self.head = SizeInvariantTimeSformer()
+
+    def forward(self, x):
+        b, t, _, h, w = x.shape
+        images = x.view(b * t, 3, h, w)
+        sequence_output = self.clip_model.encode_image(images)
+        _, _, c = sequence_output.shape
+        sequence_output = sequence_output.view(b, t, 14, 14, c)
+        sequence_output = sequence_output.permute(0, 1, 4, 2, 3)
+
+
+        res = self.head(sequence_output)
+
+        return res
+
+
+if __name__ == '__main__':
+
+    model = ViT_B_MINTIME()
+    model = model.cuda()
+    dummy_input = torch.randn(4,8,3,224,224)
+    dummy_input = dummy_input.cuda()
+    print(model(dummy_input))
+

models/NPR.py

@@ -0,0 +1,284 @@
+"""
+NPR: Rethinking the Up-Sampling Operations in CNN-based Generative Network for Generalizable Deepfake Detection@CVPR'2024
+Copyright (c) Beijing Jiaotong University and its affiliates.
+Modified by Chuangchuang Tan from https://github.com/chuangchuangtan/NPR-DeepfakeDetection
+"""
+
+import torch.nn as nn
+import torch
+import torch.utils.model_zoo as model_zoo
+from torch.nn import functional as F
+from typing import Any, cast, Dict, List, Optional, Union
+import numpy as np
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1, zero_init_residual=False):
+        super(ResNet, self).__init__()
+
+        self.unfoldSize = 2
+        self.unfoldIndex = 0
+        assert self.unfoldSize > 1
+        assert -1 < self.unfoldIndex and self.unfoldIndex < self.unfoldSize*self.unfoldSize
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64 , layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        # self.fc1 = nn.Linear(512 * block.expansion, 1)
+        self.fc1 = nn.Linear(512, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+    def interpolate(self, img, factor):
+        return F.interpolate(F.interpolate(img, scale_factor=factor, mode='nearest', recompute_scale_factor=True), scale_factor=1/factor, mode='nearest', recompute_scale_factor=True)
+    def forward(self, x):
+        # n,c,w,h = x.shape
+        # if -1*w%2 != 0: x = x[:,:,:w%2*-1,:      ]
+        # if -1*h%2 != 0: x = x[:,:,:      ,:h%2*-1]
+        # factor = 0.5
+        # x_half = F.interpolate(x, scale_factor=factor, mode='nearest', recompute_scale_factor=True)
+        # x_re   = F.interpolate(x_half, scale_factor=1/factor, mode='nearest', recompute_scale_factor=True)
+        # NPR  = x - x_re
+        # n,c,w,h = x.shape
+        # if w%2 == 1 : x = x[:,:,:-1,:]
+        # if h%2 == 1 : x = x[:,:,:,:-1]
+        b, t, _, h, w = x.shape
+        x = x.view(b * t, 3, h, w)
+        
+
+        NPR  = x - self.interpolate(x, 0.5)
+
+        x = self.conv1(NPR*2.0/3.0)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc1(x)
+        x = x.view(b, t, -1)
+        x = x.mean(1)
+
+        return x
+
+    def infer(self, x):
+        # n,c,w,h = x.shape
+        # if -1*w%2 != 0: x = x[:,:,:w%2*-1,:      ]
+        # if -1*h%2 != 0: x = x[:,:,:      ,:h%2*-1]
+        # factor = 0.5
+        # x_half = F.interpolate(x, scale_factor=factor, mode='nearest', recompute_scale_factor=True)
+        # x_re   = F.interpolate(x_half, scale_factor=1/factor, mode='nearest', recompute_scale_factor=True)
+        # NPR  = x - x_re
+        # n,c,w,h = x.shape
+        # if w%2 == 1 : x = x[:,:,:-1,:]
+        # if h%2 == 1 : x = x[:,:,:,:-1]
+        b, t, _, h, w = x.shape
+        x = x.view(b * t, 3, h, w)
+        
+
+        NPR  = x - self.interpolate(x, 0.5)
+
+        x = self.conv1(NPR*2.0/3.0)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc1(x)
+        x = x.view(b, t, -1)
+        x = x.mean(1)
+        return x
+
+
+def resnet18_npr(pretrained=False, **kwargs):
+    """Constructs a ResNet-18 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
+    return model
+
+
+def resnet34_npr(pretrained=False, **kwargs):
+    """Constructs a ResNet-34 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+    return model
+
+
+def resnet50_npr(pretrained=False, **kwargs):
+    """Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model_state = torch.load('/ossfs/workspace/aigc_video/weights/resnet50-19c8e357.pth')
+        model.load_state_dict(model_state, strict=False)
+    return model
+
+
+def resnet101_npr(pretrained=False, **kwargs):
+    """Constructs a ResNet-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+    return model
+
+
+def resnet152_npr(pretrained=False, **kwargs):
+    """Constructs a ResNet-152 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+    return model
+
+
+

models/STIL.py

@@ -0,0 +1,641 @@
+'''
+STIL: Spatiotemporal inconsistency learning for deepfake video detection @ ACM MM'2021
+Copyright (c) Tencent Youtu Lab and its affiliates.
+Modified by Zhiyuan Yan from https://github.com/Tencent/TFace?tab=readme-ov-file
+'''
+
+import os
+import datetime
+import logging
+import numpy as np
+from sklearn import metrics
+from typing import Union
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.model_zoo as model_zoo
+from torch.nn import DataParallel
+from torch.utils.tensorboard import SummaryWriter
+
+class ISM_Module(nn.Module):
+    def __init__(self, k_size=3):
+        """The Information Supplement Module (ISM).
+
+        Args:
+            k_size (int, optional): Conv1d kernel_size . Defaults to 3.
+        """
+        super(ISM_Module, self).__init__()
+
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size-1)//2, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.tensor): Input tensor of shape (nt, c, h, w)
+        """
+        y = self.avg_pool(x)
+        y = self.conv(y.squeeze(-1).transpose(-1,-2)).transpose(-1,-2).unsqueeze(-1)
+        y = self.sigmoid(y)
+        return x * y.expand_as(x)
+
+
+class TIM_Module(nn.Module):
+    def __init__(self, in_channels, reduction=16, n_segment=8, return_attn=False):
+        """The Temporal Inconsistency Module (TIM).
+
+        Args:
+            in_channels (int): Input channel number.
+            reduction (int, optional): Channel compression ratio r in the split operation.. Defaults to 16.
+            n_segment (int, optional): Number of input frames.. Defaults to 8.
+            return_attn (bool, optional): Whether to return the attention part. Defaults to False.
+
+        """
+        super(TIM_Module, self).__init__()
+        self.in_channels = in_channels
+        self.reduction = reduction
+        self.n_segment = n_segment
+        self.return_attn = return_attn
+
+        self.reduced_channels = self.in_channels // self.reduction
+
+        # first conv to shrink input channels
+        self.conv1 = nn.Conv2d(self.in_channels, self.reduced_channels, kernel_size=1, padding=0, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.reduced_channels)
+
+        self.conv_ht = nn.Conv2d(self.reduced_channels, self.reduced_channels, 
+            kernel_size=(3, 1), padding=(1, 0), groups=self.reduced_channels, bias=False)
+        self.conv_tw = nn.Conv2d(self.reduced_channels, self.reduced_channels, 
+            kernel_size=(1, 3), padding=(0, 1), groups=self.reduced_channels, bias=False)
+
+        self.avg_pool_ht = nn.AvgPool2d((2, 1), (2, 1))
+        self.avg_pool_tw = nn.AvgPool2d((1, 2), (1, 2))
+
+        # HTIE in two directions
+        self.htie_conv1 = nn.Sequential(
+            nn.Conv2d(self.reduced_channels, self.reduced_channels, kernel_size=(3, 1), padding=(1, 0), bias=False),
+            nn.BatchNorm2d(self.reduced_channels),
+        )
+        self.vtie_conv1 = nn.Sequential(
+            nn.Conv2d(self.reduced_channels, self.reduced_channels, kernel_size=(1, 3), padding=(0, 1), bias=False),
+            nn.BatchNorm2d(self.reduced_channels),
+        )
+        self.htie_conv2 = nn.Sequential(
+            nn.Conv2d(self.reduced_channels, self.reduced_channels, kernel_size=(3, 1), padding=(1, 0), bias=False),
+            nn.BatchNorm2d(self.reduced_channels),
+        )
+        self.vtie_conv2 = nn.Sequential(
+            nn.Conv2d(self.reduced_channels, self.reduced_channels, kernel_size=(1, 3), padding=(0, 1), bias=False),
+            nn.BatchNorm2d(self.reduced_channels),
+        )
+        self.ht_up_conv = nn.Sequential(
+            nn.Conv2d(self.reduced_channels, self.in_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(self.in_channels)
+        )
+        self.tw_up_conv = nn.Sequential(
+            nn.Conv2d(self.reduced_channels, self.in_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(self.in_channels)
+        )
+
+        self.sigmoid = nn.Sigmoid()
+
+
+    def feat_ht(self, feat):
+        """The H-T branch in the TIM module.
+
+        Args:
+            feat (torch.tensor): Input feature with shape [n, t, c, h, w] (c is in_channels // reduction)
+
+        """
+        n, t, c, h, w = feat.size()
+        # [n, t, c, h, w] -> [n, w, c, h, t] -> [nw, c, h, t]
+        feat_h = feat.permute(0, 4, 2, 3, 1).contiguous().view(-1, c, h, t)
+
+        # [nw, c, h, t-1]
+        feat_h_fwd, _ = feat_h.split([self.n_segment-1, 1], dim=3)
+        feat_h_conv = self.conv_ht(feat_h)
+        _, feat_h_conv_fwd = feat_h_conv.split([1, self.n_segment-1], dim=3)
+
+        diff_feat_fwd = feat_h_conv_fwd - feat_h_fwd
+        diff_feat_fwd = F.pad(diff_feat_fwd, [0, 1], value=0)  # [nw, c, h, t]
+
+        # HTIE, down_up branch
+        diff_feat_fwd1 = self.avg_pool_ht(diff_feat_fwd)  # [nw, c, h//2, t]
+        diff_feat_fwd1 = self.htie_conv1(diff_feat_fwd1)  # [nw, c, h//2, t]
+        diff_feat_fwd1 = F.interpolate(diff_feat_fwd1, diff_feat_fwd.size()[2:]) # [nw, c, h, t]
+        # HTIE, direct conv branch
+        diff_feat_fwd2 = self.htie_conv2(diff_feat_fwd)  # [nw, c, h, t]
+        
+        # [nw, C, h, t]
+        feat_ht_out = self.ht_up_conv(1/3. * diff_feat_fwd + 1/3. * diff_feat_fwd1 + 1/3. * diff_feat_fwd2)
+        feat_ht_out = self.sigmoid(feat_ht_out) - 0.5
+        # [nw, C, h, t] -> [n, w, C, h, t] -> [n, t, C, h, w]
+        feat_ht_out = feat_ht_out.view(n, w, self.in_channels, h, t).permute(0, 4, 2, 3, 1).contiguous()
+        # [n, t, C, h, w] -> [nt, C, h, w]
+        feat_ht_out = feat_ht_out.view(-1, self.in_channels, h, w)
+
+        return feat_ht_out
+
+
+    def feat_tw(self, feat):
+        """The T-W branch in the TIM module.
+
+        Args:
+            feat (torch.tensor): Input feature with shape [n, t, c, h, w] (c is in_channels // reduction)
+        """
+        n, t, c, h, w = feat.size()
+        # [n, t, c, h, w] -> [n, h, c, t, w] -> [nh, c, t, w]
+        feat_w = feat.permute(0, 3, 2, 1, 4).contiguous().view(-1, c, t, w)
+
+        # [nh, c, t-1, w]
+        feat_w_fwd, _ = feat_w.split([self.n_segment-1, 1], dim=2)
+        feat_w_conv = self.conv_tw(feat_w)
+        _, feat_w_conv_fwd = feat_w_conv.split([1, self.n_segment-1], dim=2)
+
+        diff_feat_fwd = feat_w_conv_fwd - feat_w_fwd
+        diff_feat_fwd = F.pad(diff_feat_fwd, [0, 0, 0, 1], value=0)  # [nh, c, t, w]
+
+        # VTIE, down_up branch
+        diff_feat_fwd1 = self.avg_pool_tw(diff_feat_fwd)  # [nh, c, t, w//2]
+        diff_feat_fwd1 = self.vtie_conv1(diff_feat_fwd1)  # [nh, c, t, w//2]
+        diff_feat_fwd1 = F.interpolate(diff_feat_fwd1, diff_feat_fwd.size()[2:]) # [nh, c, t, w]
+        # VTIE, direct conv branch
+        diff_feat_fwd2 = self.vtie_conv2(diff_feat_fwd)  # [nh, c, t, w]
+
+        # [nh, C, t, w]
+        feat_tw_out = self.tw_up_conv(1/3. * diff_feat_fwd + 1/3. * diff_feat_fwd1 + 1/3. * diff_feat_fwd2)
+        feat_tw_out = self.sigmoid(feat_tw_out) - 0.5
+        # [nh, C, t, w] -> [n, h, C, t, w] -> [n, t, C, h, W]
+        feat_tw_out = feat_tw_out.view(n, h, self.in_channels, t, w).permute(0, 3, 2, 1, 4).contiguous()
+        # [n, t, C, h, w] -> [nt, C, h, w]
+        feat_tw_out = feat_tw_out.view(-1, self.in_channels, h, w)
+
+        return feat_tw_out
+
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.tensor): Input with shape [nt, c, h, w]
+        """
+        # [nt, c, h, w] -> [nt, c//r, h, w]
+        bottleneck = self.conv1(x)
+        bottleneck = self.bn1(bottleneck)
+        # [nt, c//r, h, w] -> [n, t, c//r, h, w]
+        bottleneck = bottleneck.view((-1, self.n_segment) + bottleneck.size()[1:])
+
+        F_h = self.feat_ht(bottleneck)  # [nt, c, h, w]
+        F_w = self.feat_tw(bottleneck)  # [nt, c, h, w]
+
+        att = 0.5 * (F_h + F_w)
+
+        if self.return_attn:
+            return att
+
+        y2 = x + x * att
+
+        return y2
+
+
+class ShiftModule(nn.Module):
+    def __init__(self, input_channels, n_segment=8, n_div=8, mode='shift'):
+        """A depth-wise conv on the segment level.
+
+        Args:
+            input_channels (int): Input channel number.
+            n_segment (int, optional): Number of input frames.. Defaults to 8.
+            n_div (int, optional): How many channels to group as a fold.. Defaults to 8.
+            mode (str, optional): One of "shift", "fixed", "norm". Defaults to 'shift'.
+        """
+        super(ShiftModule, self).__init__()
+        self.input_channels = input_channels
+        self.n_segment = n_segment
+        self.fold_div = n_div
+        self.fold = self.input_channels // self.fold_div
+        self.conv = nn.Conv1d(self.fold_div*self.fold, self.fold_div*self.fold,
+                kernel_size=3, padding=1, groups=self.fold_div*self.fold,
+                bias=False)
+
+        if mode == 'shift':
+            self.conv.weight.requires_grad = True
+            self.conv.weight.data.zero_()
+            # shift left
+            self.conv.weight.data[:self.fold, 0, 2] = 1 
+            # shift right
+            self.conv.weight.data[self.fold: 2 * self.fold, 0, 0] = 1 
+            if 2*self.fold < self.input_channels:
+                self.conv.weight.data[2 * self.fold:, 0, 1] = 1 # fixed
+        elif mode == 'fixed':
+            self.conv.weight.requires_grad = True
+            self.conv.weight.data.zero_()
+            self.conv.weight.data[:, 0, 1] = 1 # fixed
+        elif mode == 'norm':
+            self.conv.weight.requires_grad = True
+
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.tensor): Input with shape [nt, c, h, w]
+        """
+        nt, c, h, w = x.size()
+        n_batch = nt // self.n_segment
+        x = x.view(n_batch, self.n_segment, c, h, w)
+        # (n, h, w, c, t)
+        x = x.permute(0, 3, 4, 2, 1) 
+        x = x.contiguous().view(n_batch*h*w, c, self.n_segment)
+        # (n*h*w, c, t)
+        x = self.conv(x) 
+        x = x.view(n_batch, h, w, c, self.n_segment)
+        # (n, t, c, h, w)
+        x = x.permute(0, 4, 3, 1, 2) 
+        x = x.contiguous().view(nt, c, h, w)
+        return x
+
+
+class SCConv(nn.Module):
+    """
+    The spatial conv in SIM. Used in SCBottleneck
+    """
+    def __init__(self, inplanes, planes, stride, padding, dilation, groups, pooling_r, norm_layer):
+        super(SCConv, self).__init__()
+        self.f_w = nn.Sequential(
+                    nn.AvgPool2d(kernel_size=pooling_r, stride=pooling_r), 
+                    nn.Conv2d(inplanes, planes, kernel_size=(1,3), stride=1,
+                                padding=(0,padding), dilation=(1,dilation),
+                                groups=groups, bias=False),
+                    norm_layer(planes), nn.ReLU(inplace=True))
+        self.f_h = nn.Sequential(
+                    # nn.AvgPool2d(kernel_size=(pooling_r,1), stride=(pooling_r,1)), 
+                    nn.Conv2d(inplanes, planes, kernel_size=(3,1), stride=1,
+                                padding=(padding,0), dilation=(dilation,1),
+                                groups=groups, bias=False),
+                    norm_layer(planes),
+                    )
+        self.k3 = nn.Sequential(
+                    nn.Conv2d(inplanes, planes, kernel_size=3, stride=1,
+                                padding=padding, dilation=dilation,
+                                groups=groups, bias=False),
+                    norm_layer(planes),
+                    )
+        self.k4 = nn.Sequential(
+                    nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride,
+                                padding=padding, dilation=dilation,
+                                groups=groups, bias=False),
+                    norm_layer(planes),
+                    )
+
+
+    def forward(self, x):
+        identity = x
+
+        # sigmoid(identity + k2)
+        out = torch.sigmoid(
+            torch.add(
+                identity, 
+                F.interpolate(self.f_h(self.f_w(x)), identity.size()[2:])
+            )
+        ) 
+        out = torch.mul(self.k3(x), out) # k3 * sigmoid(identity + k2)
+        s2t_info = out
+        out = self.k4(out) # k4
+
+        return out, s2t_info
+
+
+class SCBottleneck(nn.Module):
+    """
+    SCNet SCBottleneck. Variant for ResNet Bottlenect.
+    """
+    expansion = 4
+    pooling_r = 4 # down-sampling rate of the avg pooling layer in the K3 path of SC-Conv.
+
+    def __init__(self, num_segments, inplanes, planes, stride=1, downsample=None,
+                 cardinality=1, bottleneck_width=32,
+                 avd=False, dilation=1, is_first=False,
+                 norm_layer=None):
+        super(SCBottleneck, self).__init__()
+        group_width = int(planes * (bottleneck_width / 64.)) * cardinality
+        self.conv1_a = nn.Conv2d(inplanes, group_width, kernel_size=1, bias=False)
+        self.bn1_a = norm_layer(group_width)
+        self.conv1_b = nn.Conv2d(inplanes, group_width, kernel_size=1, bias=False)
+        self.bn1_b = norm_layer(group_width)
+        self.avd = avd and (stride > 1 or is_first)
+        self.tim = TIM_Module(group_width, n_segment=num_segments) 
+        self.shift = ShiftModule(group_width, n_segment=num_segments, n_div=8, mode='shift')
+        self.inplanes = inplanes
+        self.planes = planes
+        self.ism = ISM_Module()
+        self.shift = ShiftModule(group_width, n_segment=num_segments, n_div=8, mode='shift')
+
+        if self.avd:
+            self.avd_layer = nn.AvgPool2d(3, stride, padding=1)
+            stride = 1
+
+        self.k1 = nn.Sequential(
+                    nn.Conv2d(
+                        group_width, group_width, kernel_size=3, stride=stride,
+                        padding=dilation, dilation=dilation,
+                        groups=cardinality, bias=False),
+                    norm_layer(group_width),
+                    )
+
+        self.scconv = SCConv(
+            group_width, group_width, stride=stride,
+            padding=dilation, dilation=dilation,
+            groups=cardinality, pooling_r=self.pooling_r, norm_layer=norm_layer)
+
+        self.conv3 = nn.Conv2d(
+            group_width * 2, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = norm_layer(planes*4)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.dilation = dilation
+
+
+    def forward(self, x):
+        """Forward func which splits the input into two branchs a and b.
+        a: trace features
+        b: spatial features
+        """
+        residual = x
+
+        out_a = self.relu(self.bn1_a(self.conv1_a(x)))
+        out_b = self.relu(self.bn1_b(self.conv1_b(x)))
+
+        # spatial representations
+        out_b, s2t_info = self.scconv(out_b)
+        out_b = self.relu(out_b)
+
+        # trace features
+        out_a = self.tim(out_a)
+        out_a = self.shift(out_a + self.ism(s2t_info))
+        out_a = self.relu(self.k1(out_a))
+
+        if self.avd:
+            out_a = self.avd_layer(out_a)
+            out_b = self.avd_layer(out_b)
+
+        out = self.conv3(torch.cat([out_a, out_b], dim=1))
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class SCNet(nn.Module):
+    def __init__(self, num_segments, block, layers, groups=1, bottleneck_width=32,
+                 num_classes=1000, dilated=False, dilation=1,
+                 deep_stem=False, stem_width=64, avg_down=False,
+                 avd=False, norm_layer=nn.BatchNorm2d):
+        """SCNet, a variant based on ResNet.
+
+        Args:
+            num_segments (int): 
+                Number of input frames.
+            block (class): 
+                Class for the residual block.
+            layers (list): 
+                Number of layers in each block.
+            num_classes (int, optional): 
+                Number of classification class.. Defaults to 1000.
+            dilated (bool, optional): 
+                Whether to apply dilation conv. Defaults to False.
+            dilation (int, optional): 
+                The dilation parameter in dilation conv. Defaults to 1.
+            deep_stem (bool, optional): 
+                Whether to replace 7x7 conv in input stem with 3 3x3 conv. Defaults to False.
+            stem_width (int, optional): 
+                Stem width in conv1 stem. Defaults to 64.
+            avg_down (bool, optional): 
+                Whether to use AvgPool instead of stride conv when downsampling in the bottleneck. Defaults to False.
+            avd (bool, optional): 
+                The avd parameter for the block Defaults to False.
+            norm_layer (class, optional): 
+                Normalization layer. Defaults to nn.BatchNorm2d.
+        """
+        self.cardinality = groups
+        self.bottleneck_width = bottleneck_width
+        # ResNet-D params
+        self.inplanes = stem_width*2 if deep_stem else 64
+        self.avg_down = avg_down
+        self.avd = avd
+        self.num_segments = num_segments
+
+        super(SCNet, self).__init__()
+        conv_layer = nn.Conv2d
+        if deep_stem:
+            self.conv1 = nn.Sequential(
+                conv_layer(3, stem_width, kernel_size=3, stride=2, padding=1, bias=False),
+                norm_layer(stem_width),
+                nn.ReLU(inplace=True),
+                conv_layer(stem_width, stem_width, kernel_size=3, stride=1, padding=1, bias=False),
+                norm_layer(stem_width),
+                nn.ReLU(inplace=True),
+                conv_layer(stem_width, stem_width*2, kernel_size=3, stride=1, padding=1, bias=False),
+            )
+        else:
+            self.conv1 = conv_layer(3, 64, kernel_size=7, stride=2, padding=3,
+                                   bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer, is_first=False)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer)
+        if dilated or dilation == 4:
+            self.layer3 = self._make_layer(block, 256, layers[2], stride=1,
+                                           dilation=2, norm_layer=norm_layer)
+            self.layer4 = self._make_layer(block, 512, layers[3], stride=1,
+                                           dilation=4, norm_layer=norm_layer)
+        elif dilation==2:
+            self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                           dilation=1, norm_layer=norm_layer)
+            self.layer4 = self._make_layer(block, 512, layers[3], stride=1,
+                                           dilation=2, norm_layer=norm_layer)
+        else:
+            self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                           norm_layer=norm_layer)
+            self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                           norm_layer=norm_layer)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, norm_layer):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=None,
+                    is_first=True):
+        """
+        Core function to build layers.
+        """
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            down_layers = []
+            if self.avg_down:
+                if dilation == 1:
+                    down_layers.append(nn.AvgPool2d(kernel_size=stride, stride=stride,
+                                                    ceil_mode=True, count_include_pad=False))
+                else:
+                    down_layers.append(nn.AvgPool2d(kernel_size=1, stride=1,
+                                                    ceil_mode=True, count_include_pad=False))
+                down_layers.append(nn.Conv2d(self.inplanes, planes * block.expansion,
+                                             kernel_size=1, stride=1, bias=False))
+            else:
+                down_layers.append(nn.Conv2d(self.inplanes, planes * block.expansion,
+                                             kernel_size=1, stride=stride, bias=False))
+            down_layers.append(norm_layer(planes * block.expansion))
+            downsample = nn.Sequential(*down_layers)
+
+        layers = []
+        if dilation == 1 or dilation == 2:
+            layers.append(block(self.num_segments, self.inplanes, planes, stride, downsample=downsample,
+                                cardinality=self.cardinality,
+                                bottleneck_width=self.bottleneck_width,
+                                avd=self.avd, dilation=1, is_first=is_first, 
+                                norm_layer=norm_layer))
+        elif dilation == 4:
+            layers.append(block(self.num_segments, self.inplanes, planes, stride, downsample=downsample,
+                                cardinality=self.cardinality,
+                                bottleneck_width=self.bottleneck_width,
+                                avd=self.avd, dilation=2, is_first=is_first, 
+                                norm_layer=norm_layer))
+        else:
+            raise RuntimeError("=> unknown dilation size: {}".format(dilation))
+
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.num_segments, self.inplanes, planes,
+                                cardinality=self.cardinality,
+                                bottleneck_width=self.bottleneck_width,
+                                avd=self.avd, dilation=dilation, 
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+
+    def features(self, input):
+        x = self.conv1(input)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+
+    def logits(self, features):
+        x = self.avgpool(features)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+
+    def forward(self, input):
+        x = self.features(input)
+        x = self.logits(x)
+        return x
+
+
+def scnet50_v1d(num_segments, pretrained=False, **kwargs):
+    """
+    SCNet backbone, which is based on ResNet-50
+    Args:
+        num_segments (int):
+            Number of input frames.
+        pretrained (bool, optional):
+            Whether to load pretrained weights.
+    """
+    model = SCNet(num_segments, SCBottleneck, [3, 4, 6, 3],
+                   deep_stem=True, stem_width=32, avg_down=True,
+                   avd=True, **kwargs)
+    if pretrained:
+        model_state = torch.load('/ossfs/workspace/GenVideo/pretrained_weights/scnet50_v1d-4109d1e1.pth')
+        model.load_state_dict(model_state, strict=False)
+
+    return model
+
+class STIL_Model(nn.Module):
+    def __init__(self, 
+        num_class=1,
+        num_segment=8,
+        add_softmax=False,
+        **kwargs):
+        """ Model Builder for STIL model.
+        STIL: Spatiotemporal Inconsistency Learning for DeepFake Video Detection (https://arxiv.org/abs/2109.01860)
+        
+        Args:
+            num_class (int, optional): Number of classes. Defaults to 2.
+            num_segment (int, optional): Number of segments (frames) fed to the model. Defaults to 8.
+            add_softmax (bool, optional): Whether to add softmax layer at the end. Defaults to False.
+        """
+        super().__init__()
+
+        self.num_class = num_class
+        self.num_segment = num_segment
+        self.build_model()
+
+    def build_model(self):
+        self.base_model = scnet50_v1d(self.num_segment, pretrained=True)
+
+        fc_feature_dim = self.base_model.fc.in_features
+        self.base_model.fc = nn.Linear(fc_feature_dim, self.num_class)
+
+    def forward(self, x):
+        """Forward pass of the model.
+        
+        Args:
+            x (torch.tensor): input tensor of shape (n, t*c, h, w). n is the batch_size, t is num_segment
+        """
+        # img channel default to 3
+        img_channel = 3
+
+        # x: [n, tc, h, w] -> [nt, c, h, w]
+        # out: [nt, num_class]
+        out = self.base_model(
+            x.view((-1, img_channel) + x.size()[2:])
+        )
+
+        out = out.view(-1, self.num_segment, self.num_class)  # [n, t, num_class]
+        out = out.mean(1, keepdim=False)  # [n, num_class]
+        return out
+
+
+    def set_segment(self, num_segment):
+        """Change num_segment of the model. 
+        Useful when the train and test want to feed different number of frames.
+
+        Args:
+            num_segment (int): New number of segments.
+        """
+        self.num_segment = num_segment
+
+
+class Det_STIL(nn.Module):
+    def __init__(self):
+        super(Det_STIL, self).__init__()
+        self.model = STIL_Model()
+
+    def forward(self, x):
+        b, t, _, h, w = x.shape
+        images = x.view(b, t*3, h, w)
+        x = self.model(images)
+        return x

models/TALL.py

@@ -0,0 +1,935 @@
+# --------------------------------------------------------
+# Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.models.registry import register_model
+from torch.hub import load_state_dict_from_url
+import logging
+from einops import rearrange
+import math
+
+_logger = logging.getLogger(__name__)
+
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
+        'first_conv': 'patch_embed.proj', 'classifier': 'head',
+        **kwargs
+    }
+
+
+default_cfgs = {
+    # patch models (my experiments)
+    'swin_base_in1k_patch4_224': _cfg(
+        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224.pth',
+    ),
+    'swin_base_patch4_window7_224_22k': _cfg(
+        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth',
+    )
+}
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+# adjust image size for the pyramid structure (i.e. must be integer of 32)
+def create_new_image_size(img_size, thumbnail_dim, window_size):
+    h, w = img_size * thumbnail_dim[0], img_size * thumbnail_dim[1]
+    new_h, new_w = h, w
+    dim = 32 * window_size
+    if h % (32 * window_size) != 0:
+        new_h = (h // dim + 1) * dim
+    if w % (32 * window_size) != 0:
+        new_w = (w // dim + 1) * dim
+
+    return (new_h//thumbnail_dim[0], new_w//thumbnail_dim[1])
+
+
+class WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, bottleneck=False, use_checkpoint=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        self.use_checkpoint = use_checkpoint
+
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            w_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = cnt
+                    cnt += 1
+
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward_attn(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+        x = x.view(B, H * W, C)
+
+        return x
+
+    def forward_mlp(self, x):
+        return self.drop_path(self.mlp(self.norm2(x)))
+
+    def forward(self, x):
+        shortcut = x
+        if self.use_checkpoint:
+            x = checkpoint.checkpoint(self.forward_attn, x)
+        else:
+            x = self.forward_attn(x)
+        x = shortcut + self.drop_path(x)
+
+        if self.use_checkpoint:
+            x = x + checkpoint.checkpoint(self.forward_mlp, x)
+        else:
+            x = x + self.forward_mlp(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        nW = H * W / self.window_size / self.window_size
+        flops += nW * self.attn.flops(self.window_size * self.window_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        return flops
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 bottleneck=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, window_size=window_size,
+                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer,
+                                 bottleneck=bottleneck if i == depth-1 else False,
+                                 use_checkpoint=use_checkpoint)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+class PatchEmbed(nn.Module):
+    r""" Image to Patch Embedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=(224,224), patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        #img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+    def flops(self):
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class SwinTransformer(nn.Module):
+    r""" Swin Transformer
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+    """
+
+    def __init__(self, duration=8, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
+                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                 use_checkpoint=False, thumbnail_rows=1, bottleneck=False, **kwargs):
+        super().__init__()
+
+        self.duration = duration
+        self.num_classes = num_classes
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+        self.thumbnail_rows = thumbnail_rows
+
+        self.img_size = img_size
+        self.window_size = [window_size for _ in depths] if not isinstance(window_size, list) else window_size
+        self.image_mode = True
+
+        self.frame_padding = self.duration % thumbnail_rows if self.image_mode is True else 0
+        if self.frame_padding != 0:
+            self.frame_padding = self.thumbnail_rows - self.frame_padding
+            self.duration += self.frame_padding
+
+        # split image into non-overlapping patches
+        if self.image_mode:
+            thumbnail_dim = (thumbnail_rows, self.duration // thumbnail_rows)
+            thumbnail_size = (img_size * thumbnail_dim[0], img_size * thumbnail_dim[1])
+        else:
+           thumbnail_size = (img_size, img_size)
+        
+        print ('---------------------------------------')
+        print ('duration:', self.duration, 'frame padding:', self.frame_padding, 'image_size:', self.img_size, 'patch_size:', patch_size, 'thumbnail_size:', (thumbnail_rows, self.duration//thumbnail_rows), 'ape:', self.ape)
+
+        self.patch_embed = PatchEmbed(
+            img_size=(img_size, img_size), patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # absolute position embedding
+        if self.ape:
+#            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+#            trunc_normal_(self.absolute_pos_embed, std=.02)
+            self.frame_pos_embed = nn.Parameter(torch.zeros(1, self.duration, embed_dim))
+            trunc_normal_(self.frame_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=depths[i_layer],
+                               num_heads=num_heads[i_layer],
+                               window_size=self.window_size[i_layer],
+                               mlp_ratio=self.mlp_ratio,
+                               qkv_bias=qkv_bias, qk_scale=qk_scale,
+                               drop=drop_rate, attn_drop=attn_drop_rate,
+                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                               norm_layer=norm_layer,
+                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                               use_checkpoint=use_checkpoint,
+                               bottleneck=bottleneck)
+            self.layers.append(layer)
+
+        self.norm = norm_layer(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed', 'frame_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def create_thumbnail(self, x):
+        # import pdb;pdb.set_trace()
+        input_size = x.shape[-2:]
+        if input_size != to_2tuple(self.img_size):
+            x = nn.functional.interpolate(x, size=self.img_size,mode='bilinear')
+        x = rearrange(x, 'b (th tw c) h w -> b c (th h) (tw w)', th=self.thumbnail_rows, c=3)
+        return x
+
+    def pad_frames(self, x):
+        frame_num = self.duration - self.frame_padding
+        x = x.view((-1,3*frame_num)+x.size()[2:])
+        x_padding = torch.zeros((x.shape[0], 3*self.frame_padding) + x.size()[2:]).cuda()
+        x = torch.cat((x, x_padding), dim=1)
+        assert x.shape[1] == 3 * self.duration, 'frame number %d not the same as adjusted input size %d' % (x.shape[1], 3 * self.duration)
+
+        return x
+
+    # need to find a better way to do this, maybe torch.fold?
+    def create_image_pos_embed(self):
+        img_rows, img_cols = self.patches_resolution
+        _, _, T = self.frame_pos_embed.shape
+        rows = img_rows // self.thumbnail_rows
+        cols = img_cols // (self.duration // self.thumbnail_rows)
+        img_pos_embed = torch.zeros(img_rows, img_cols, T).cuda()
+         #print (self.duration, T, img_rows, img_cols, rows, cols)
+        for i in range(self.duration):
+            r_indx = (i // self.thumbnail_rows) * rows
+            c_indx = (i % self.thumbnail_rows) * cols
+            img_pos_embed[r_indx:r_indx+rows,c_indx:c_indx+cols] = self.frame_pos_embed[0, i]
+            #print (r_indx, r_indx+rows, c_indx, c_indx+cols)
+        return img_pos_embed.reshape(-1, T)
+
+    def forward_features(self, x):
+#        x = rearrange(x, 'b (t c) h w -> b c h (t w)', t=self.duration)
+        # in evaluation, it's Bx(num_crops*num_cips*num_frames*3)xHxW
+        # import pdb;pdb.set_trace()
+        b, t, _, h, w = x.shape
+        x = x.view(b, t*3, h, w)
+        if self.frame_padding > 0:
+            x = self.pad_frames(x)
+        else:
+            x = x.view((-1,3*self.duration)+x.size()[2:])
+
+        if self.image_mode:
+            x = self.create_thumbnail(x) 
+            x = nn.functional.interpolate(x, size=self.img_size,mode='bilinear')
+        else:
+            x = rearrange(x, 'b (n t c) h w -> (b n t) c h w', t=self.duration, c=3)
+       
+        x = self.patch_embed(x)
+        if self.ape:
+#           x = x + self.absolute_pos_embed
+            img_pos_embed = self.create_image_pos_embed()
+            x = x + img_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = self.norm(x)  # B L C
+        x = self.avgpool(x.transpose(1, 2))  # B C 1
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        if not self.image_mode:
+            x = x.view(-1, self.duration, self.num_classes)
+            x = torch.mean(x, dim=1) 
+        return x
+
+    def flops(self):
+        flops = 0
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            flops += layer.flops()
+        flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
+        flops += self.num_features * self.num_classes
+        return flops
+
+
+def load_pretrained(model, cfg=None, num_classes=1000, in_chans=3, filter_fn=None, img_size=224, num_patches=196,
+                 pretrained_window_size=7, pretrained_model="", strict=True):
+    if cfg is None:
+        cfg = getattr(model, 'default_cfg')
+    if cfg is None or 'url' not in cfg or not cfg['url']:
+        _logger.warning("Pretrained model URL is invalid, using random initialization.")
+        return
+
+    if len(pretrained_model) == 0:
+        # state_dict = model_zoo.load_url(cfg['url'], progress=False, map_location='cpu')
+        # state_dict = load_state_dict_from_url(cfg['url'], progress=False, map_location='cpu')
+        state_dict = torch.load('/mnt/new_nas/yansan/models/swin_base_patch4_window7_224_22k.pth', map_location='cpu')
+    else:
+        try:
+            state_dict = load_state_dict(pretrained_model)['model']
+        except:
+            state_dict = load_state_dict(pretrained_model)
+
+    if filter_fn is not None:
+        state_dict = filter_fn(state_dict)
+
+    if in_chans == 1:
+        conv1_name = cfg['first_conv']
+        _logger.info('Converting first conv (%s) pretrained weights from 3 to 1 channel' % conv1_name)
+        conv1_weight = state_dict[conv1_name + '.weight']
+        conv1_type = conv1_weight.dtype
+        conv1_weight = conv1_weight.float()
+        O, I, J, K = conv1_weight.shape
+        if I > 3:
+            assert conv1_weight.shape[1] % 3 == 0
+            # For models with space2depth stems
+            conv1_weight = conv1_weight.reshape(O, I // 3, 3, J, K)
+            conv1_weight = conv1_weight.sum(dim=2, keepdim=False)
+        else:
+            conv1_weight = conv1_weight.sum(dim=1, keepdim=True)
+        conv1_weight = conv1_weight.to(conv1_type)
+        state_dict[conv1_name + '.weight'] = conv1_weight
+    elif in_chans != 3:
+        conv1_name = cfg['first_conv']
+        conv1_weight = state_dict[conv1_name + '.weight']
+        conv1_type = conv1_weight.dtype
+        conv1_weight = conv1_weight.float()
+        O, I, J, K = conv1_weight.shape
+        if I != 3:
+            _logger.warning('Deleting first conv (%s) from pretrained weights.' % conv1_name)
+            del state_dict[conv1_name + '.weight']
+            strict = False
+        else:
+            _logger.info('Repeating first conv (%s) weights in channel dim.' % conv1_name)
+            repeat = int(math.ceil(in_chans / 3))
+            conv1_weight = conv1_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
+            conv1_weight *= (3 / float(in_chans))
+            conv1_weight = conv1_weight.to(conv1_type)
+            state_dict[conv1_name + '.weight'] = conv1_weight
+
+    #for key, value in state_dict['model'].items():
+    #    print (key)
+
+    classifier_name = cfg['classifier']
+    if num_classes == 1000 and cfg['num_classes'] == 1001:
+        # special case for imagenet trained models with extra background class in pretrained weights
+        classifier_weight = state_dict[classifier_name + '.weight']
+        state_dict[classifier_name + '.weight'] = classifier_weight[1:]
+        classifier_bias = state_dict[classifier_name + '.bias']
+        state_dict[classifier_name + '.bias'] = classifier_bias[1:]
+    elif num_classes != cfg['num_classes']:  # and len(pretrained_model) == 0:
+        # completely discard fully connected for all other differences between pretrained and created model
+        del state_dict['model'][classifier_name + '.weight']
+        del state_dict['model'][classifier_name + '.bias']
+        strict = False
+    '''
+    ## Resizing the positional embeddings in case they don't match
+    if img_size != cfg['input_size'][1]:
+        pos_embed = state_dict['pos_embed']
+        cls_pos_embed = pos_embed[0, 0, :].unsqueeze(0).unsqueeze(1)
+        other_pos_embed = pos_embed[0, 1:, :].unsqueeze(0).transpose(1, 2)
+        new_pos_embed = F.interpolate(other_pos_embed, size=(num_patches), mode='nearest')
+        new_pos_embed = new_pos_embed.transpose(1, 2)
+        new_pos_embed = torch.cat((cls_pos_embed, new_pos_embed), 1)
+        state_dict['pos_embed'] = new_pos_embed
+   '''
+
+    # remove window_size related parameters
+    window_size = (model.window_size)[0]
+    print (pretrained_window_size, window_size)
+
+    new_state_dict = state_dict['model'].copy()
+    for key in state_dict['model']:
+        if 'attn_mask' in key:
+            del new_state_dict[key]
+
+        #if window_size != pretrained_window_size:
+        if 1:
+            if 'relative_position_index' in key:
+                del new_state_dict[key]
+
+             # resize it
+            if 'relative_position_bias_table' in key:
+                #print ('resizing relative_position_bias_table')
+                pretrained_table = state_dict['model'][key]
+                pretrained_table_size = int(math.sqrt(pretrained_table.shape[0]))
+                table_size = int(math.sqrt(model.state_dict()[key].shape[0]))
+                #print (pretrained_table_size, table_size)
+                if pretrained_table_size != table_size:
+                    table = pretrained_table.permute(1, 0).view(1, -1, pretrained_table_size, pretrained_table_size)
+                    table = nn.functional.interpolate(table, size=table_size, mode='bilinear')
+                    table = table.view(-1, table_size*table_size).permute(1, 0)
+                    new_state_dict[key] = table
+
+    for key in model.state_dict():
+        if 'bottleneck_norm' in key:
+            attn_key = key.replace('bottleneck_norm','norm1')
+            #print (key, attn_key)
+            new_state_dict[key] = new_state_dict[attn_key]
+
+    '''
+    for key in new_state_dict:
+       if key not in model.state_dict():
+           print ('----', key)
+       else:
+           print ('++++', key)
+    print ('====================')
+    for key in model.state_dict():
+       if key not in new_state_dict:
+           print ('----', key)
+       else:
+           print ('++++', key)
+    '''
+
+    print ('loading weights....')
+    ## Loading the weights
+    model_dict = model.state_dict()
+    pretrained_dict = {k: v for k, v in new_state_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
+    model_dict.update(pretrained_dict)
+    model.load_state_dict(model_dict, strict=False)
+
+def _conv_filter(state_dict, patch_size=4):
+    """ convert patch embedding weight from manual patchify + linear proj to conv"""
+    out_dict = {}
+    for k, v in state_dict.items():
+        if 'patch_embed.proj.weight' in k:
+            if v.shape[-1] != patch_size:
+                patch_size = v.shape[-1]
+            v = v.reshape((v.shape[0], 3, patch_size, patch_size))
+        out_dict[k] = v
+    return out_dict
+
+
+def _create_vision_transformer(variant, pretrained=False, pretrained_window_size=7, **kwargs):
+    default_cfg = default_cfgs[variant]
+    default_num_classes = default_cfg['num_classes']
+    default_img_size = default_cfg['input_size'][-1]
+
+    num_classes = kwargs.pop('num_classes', default_num_classes)
+    img_size = kwargs.pop('img_size', default_img_size)
+    repr_size = kwargs.pop('representation_size', None)
+
+    model_cls = SwinTransformer
+    model = model_cls(img_size=img_size, num_classes=num_classes, **kwargs)
+    model.default_cfg = default_cfg
+
+   
+    if pretrained:
+        load_pretrained(
+            model, num_classes=num_classes, in_chans=kwargs.get('in_chans', 3),
+            filter_fn=_conv_filter,
+            img_size=img_size,
+            pretrained_window_size=pretrained_window_size,
+            pretrained_model=''
+        )
+    return model
+
+
+
+@register_model
+def TALL_SWIN(pretrained=False, **kwargs):
+    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    temporal_module_name = kwargs.pop('temporal_module_name', None)
+    temporal_attention_only = kwargs.pop('temporal_attention_only', None)
+    temporal_heads_scale = kwargs.pop('temporal_heads_scale', 1.0)
+    temporal_mlp_scale = kwargs.pop('temporal_mlp_scale', 1.0)
+    rel_pos = kwargs.pop('rel_pos', False)
+    token_maks = kwargs.pop('token_mask', False)
+    frame_cls_tokens = kwargs.pop('frame_cls_tokens', 1)
+    kwargs.pop('hub_attention', '')
+    kwargs.pop('hub_aggregation', '')
+    kwargs.pop('spatial_hub_size', (-1, -1))
+    kwargs.pop('temporal_pooling', None)
+    kwargs.pop('window_size', -1) 
+
+    embed_dim = 128
+    mlp_ratio = 4.
+    #drop_path_rate=0.5
+    patch_size=4
+    window_size=[14,14,14,7]
+    depths = [2, 2, 18, 2]
+    num_heads = [4, 8, 16, 32]
+    use_checkpoint=kwargs.pop('use_checkpoint', False)
+    ape = kwargs.pop('hpe_to_token', False)
+    bottleneck = True if kwargs.pop('bottleneck', None) is not None else False
+    model_kwargs = dict(patch_size=patch_size, window_size=window_size, embed_dim=embed_dim, depths=depths, num_heads=num_heads, mlp_ratio=mlp_ratio,
+                        use_checkpoint=use_checkpoint, ape=ape, bottleneck=bottleneck, **kwargs)
+    print(model_kwargs)
+    model = _create_vision_transformer('swin_base_patch4_window7_224_22k', pretrained=pretrained, pretrained_window_size=7, **model_kwargs)
+    return model
+
+if __name__ == '__main__':
+    dummy_input = torch.randn(4,8,3,224,224)
+    model = TALL_SWIN(pretrained=True)
+    print(model(dummy_input))

models/VideoMAE.py

@@ -0,0 +1,67 @@
+import torch
+from transformers import VideoMAEImageProcessor, VideoMAEForVideoClassification
+import torch.nn as nn
+import torchvision
+import time
+from .mamba_base import MambaConfig, ResidualBlock
+
+def create_reorder_index(N, device):
+    new_order = []
+    for col in range(N):
+        if col % 2 == 0:
+            new_order.extend(range(col, N*N, N))
+        else:
+            new_order.extend(range(col + N*(N-1), col-1, -N))
+    return torch.tensor(new_order, device=device)
+
+def reorder_data(data, N):
+    assert isinstance(data, torch.Tensor), "data should be a torch.Tensor"
+    device = data.device
+    new_order = create_reorder_index(N, device)
+    B, t, _, _ = data.shape
+    index = new_order.repeat(B, t, 1).unsqueeze(-1)
+    reordered_data = torch.gather(data, 2, index.expand_as(data))
+    return reordered_data
+
+class Videomae_Net(nn.Module):
+    def __init__(
+        self, channel_size=512, dropout=0.2, class_num=1
+    ):
+        super(Videomae_Net, self).__init__()
+        self.model = VideoMAEForVideoClassification.from_pretrained("/ossfs/workspace/GenVideo/pretrained_weights/videomae")
+        self.fc1 = nn.Linear(768, class_num)
+        self.bn1 = nn.BatchNorm1d(768)
+
+        self._init_params()
+    
+    def _init_params(self):
+        nn.init.xavier_normal_(self.fc1.weight)
+        nn.init.constant_(self.fc1.bias, 0)
+
+    def forward(self, x):
+        x = self.model.videomae(x)
+        sequence_output = x[0]
+        print(sequence_output.shape)
+        if self.model.fc_norm is not None:
+            sequence_output = self.model.fc_norm(sequence_output.mean(1))
+        else:
+            sequence_output = sequence_output[:, 0]
+        x = self.bn1(sequence_output)
+        x = self.fc1(x)
+        return x
+
+
+
+if __name__ == '__main__':
+
+    model = Videomae_Net()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    model.eval()
+
+
+    input_data = torch.randn(1, 16, 3, 224, 224).to(device)
+
+    model(input_data)
+

models/XCLIP.py

@@ -0,0 +1,33 @@
+from transformers import XCLIPVisionModel
+import os
+import sys
+import numpy as np
+import torch
+import torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+import math
+
+from transformers import XCLIPVisionModel
+class XCLIP(nn.Module):
+    def __init__(
+        self, channel_size=512, dropout=0.2, class_num=1
+    ):
+        super(XCLIP, self).__init__()
+      
+        self.backbone = XCLIPVisionModel.from_pretrained("GenVideo/pretrained_weights/xclip")
+        self.fc_norm = nn.LayerNorm(768)
+        self.head = nn.Linear(768, 1)
+
+    def forward(self, x):
+        b, t, _, h, w = x.shape
+        images = x.view(b * t, 3, h, w)
+        outputs = self.backbone(images, output_hidden_states=True)
+        sequence_output = outputs['pooler_output'].reshape(b, t, -1)
+        video_level_features = self.fc_norm(sequence_output.mean(1))
+        pred = self.head(video_level_features)
+
+        return pred
+
+

models/__init__.py

@@ -0,0 +1,5 @@
+from .F3Net import Det_F3_Net
+from .NPR import resnet50_npr
+from .STIL import Det_STIL
+from .DeMamba import XCLIP_DeMamba, CLIP_DeMamba
+from .XCLIP import XCLIP

models/clip/__init__.py

@@ -0,0 +1 @@
+from .clip import *

models/clip/clip.py

@@ -0,0 +1,233 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from tqdm import tqdm
+
+from .model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if torch.__version__.split(".") < ["1", "7", "1"]:
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    # if name in _MODELS:
+    #     model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    # elif os.path.isfile(name):
+    #     model_path = name
+    # else:
+    #     raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    model_path = 'weights/openclip/' + name +'.pt'
+    print(model_path)
+    try:
+        # loading JIT archive
+        model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
+        state_dict = None
+    except RuntimeError:
+        # loading saved state dict
+        if jit:
+            warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+            jit = False
+        state_dict = torch.load(model_path, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model, _transform(model.visual.input_resolution)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, _transform(model.input_resolution.item())
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> torch.LongTensor:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

models/clip/model.py

@@ -0,0 +1,432 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.avgpool = nn.AvgPool2d(2)
+        self.relu = nn.ReLU(inplace=True)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2), (self.conv3, self.bn3)]:
+                x = self.relu(bn(conv(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 1:, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

models/clip/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

models/mamba_base.py

@@ -0,0 +1,352 @@
+"""
+Mamba: Linear-Time Sequence Modeling with Selective State Spaces
+Copyright (c) Carnegie Mellon University.
+Implemented by alxndrTL from https://github.com/alxndrTL/mamba.py
+"""
+
+
+import math
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import DropPath
+from .pscan import pscan
+
+@dataclass
+class MambaConfig:
+    d_model: 768 # D
+    dt_rank: Union[int, str] = 'auto'
+    d_state: int = 16 # N in paper/comments
+    expand_factor: int = 2 # E in paper/comments
+    d_conv: int = 4
+
+    dt_min: float = 0.001
+    dt_max: float = 0.1
+    dt_init: str = "random" # "random" or "constant"
+    dt_scale: float = 1.0
+    dt_init_floor = 1e-4
+
+    drop_prob: float = 0.1
+
+    bias: bool = False
+    conv_bias: bool = True
+    bimamba: bool = True
+
+    pscan: bool = True # use parallel scan mode or sequential mode when training
+
+    def __post_init__(self):
+        self.d_inner = self.expand_factor * self.d_model # E*D = ED in comments
+
+        if self.dt_rank == 'auto':
+            self.dt_rank = math.ceil(self.d_model / 16)
+
+class ResidualBlock(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.mixer = MambaBlock(config)
+        self.norm = RMSNorm(config.d_model)
+        self.drop_path = DropPath(drop_prob=config.drop_prob)
+
+    def forward(self, x):
+        # x : (B, L, D)
+
+        # output : (B, L, D)
+
+        output = self.drop_path(self.mixer(self.norm(x))) + x
+        return output
+    
+    def step(self, x, cache):
+        # x : (B, D)
+        # cache : (h, inputs)
+                # h : (B, ED, N)
+                # inputs: (B, ED, d_conv-1)
+
+        # output : (B, D)
+        # cache : (h, inputs)
+
+        output, cache = self.mixer.step(self.norm(x), cache)
+        output = output + x
+        return output, cache
+
+class MambaBlock(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.config = config
+
+        # projects block input from D to 2*ED (two branches)
+        self.in_proj = nn.Linear(config.d_model, 2 * config.d_inner, bias=config.bias)
+
+        self.conv1d = nn.Conv1d(in_channels=config.d_inner, out_channels=config.d_inner, 
+                              kernel_size=config.d_conv, bias=config.conv_bias, 
+                              groups=config.d_inner,
+                              padding=config.d_conv - 1)
+        
+        # projects x to input-dependent Δ, B, C
+        self.x_proj = nn.Linear(config.d_inner, config.dt_rank + 2 * config.d_state, bias=False)
+
+        # projects Δ from dt_rank to d_inner
+        self.dt_proj = nn.Linear(config.dt_rank, config.d_inner, bias=True)
+
+        # dt initialization
+        # dt weights
+        dt_init_std = config.dt_rank**-0.5 * config.dt_scale
+        if config.dt_init == "constant":
+            nn.init.constant_(self.dt_proj.weight, dt_init_std)
+        elif config.dt_init == "random":
+            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
+        else:
+            raise NotImplementedError
+        
+        # dt bias
+        dt = torch.exp(
+            torch.rand(config.d_inner) * (math.log(config.dt_max) - math.log(config.dt_min)) + math.log(config.dt_min)
+        ).clamp(min=config.dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt)) # inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        with torch.no_grad():
+            self.dt_proj.bias.copy_(inv_dt)
+        #self.dt_proj.bias._no_reinit = True # initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
+        # todo : explain why removed
+
+        # S4D real initialization
+        A = torch.arange(1, config.d_state + 1, dtype=torch.float32).repeat(config.d_inner, 1)
+        self.A_log = nn.Parameter(torch.log(A)) # why store A in log ? to keep A < 0 (cf -torch.exp(...)) ? for gradient stability ?
+        self.D = nn.Parameter(torch.ones(config.d_inner))
+
+        # projects block output from ED back to D
+        self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=config.bias)
+
+        self.bimamba = config.bimamba
+
+        if self.bimamba:
+            A_b = torch.arange(1, config.d_state + 1, dtype=torch.float32).repeat(config.d_inner, 1)
+            self.A_b_log = nn.Parameter(torch.log(A_b))
+
+            self.conv1d_b = nn.Conv1d(in_channels=config.d_inner, out_channels=config.d_inner, 
+                              kernel_size=config.d_conv, bias=config.conv_bias, 
+                              groups=config.d_inner,
+                              padding=config.d_conv - 1)
+
+            self.x_proj_b = nn.Linear(config.d_inner, config.dt_rank + 2 * config.d_state, bias=False)
+            self.dt_proj_b = nn.Linear(config.dt_rank, config.d_inner, bias=True)
+            self.D_b = nn.Parameter(torch.ones(config.d_inner))
+
+    def forward(self, x):
+        # x : (B, L, D)
+        
+        # y : (B, L, D)
+
+        _, L, _ = x.shape
+
+        xz = self.in_proj(x) # (B, L, 2*ED)
+        x, z = xz.chunk(2, dim=-1) # (B, L, ED), (B, L, ED)
+
+        # x branch
+        x = x.transpose(1, 2) # (B, ED, L)
+        x = self.conv1d(x)[:, :, :L] # depthwise convolution over time, with a short filter
+        x = x.transpose(1, 2) # (B, L, ED)
+
+        x = F.silu(x)
+        y = self.ssm(x)
+
+        # z branch
+        z = F.silu(z)
+
+        output = y * z
+        output = self.out_proj(output) # (B, L, D)
+
+        return output
+    
+    def ssm(self, x):
+        # x : (B, L, ED)
+
+        # y : (B, L, ED)
+
+        A = -torch.exp(self.A_log.float()) # (ED, N)
+        D = self.D.float()
+        # TODO remove .float()
+
+        deltaBC = self.x_proj(x) # (B, L, dt_rank+2*N)
+
+        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, L, dt_rank), (B, L, N), (B, L, N)
+        delta = F.softplus(self.dt_proj(delta)) # (B, L, ED)
+
+        if self.config.pscan:
+            y = self.selective_scan(x, delta, A, B, C, D)
+        else:
+            y = self.selective_scan_seq(x, delta, A, B, C, D)
+
+        if self.bimamba:
+            x_b = x.flip([-1])
+            A_b = -torch.exp(self.A_b_log.float()) # (ED, N)
+            D_b = self.D_b.float()
+            deltaBC_b = self.x_proj_b(x_b)
+            delta_b, B_b, C_b = torch.split(deltaBC_b, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, L, dt_rank), (B, L, N), (B, L, N)
+            delta_b = F.softplus(self.dt_proj_b(delta_b)) # (B, L, ED)
+            if self.config.pscan:
+                y_b = self.selective_scan(x_b, delta_b, A_b, B_b, C_b, D_b)
+            else:
+                y_b = self.selective_scan_seq(x_b, delta_b, A_b, B_b, C_b, D_b)
+            y_b = y_b.flip([-1])
+            y = y + y_b
+        return y
+    
+    def selective_scan(self, x, delta, A, B, C, D):
+        # x : (B, L, ED)
+        # Δ : (B, L, ED)
+        # A : (ED, N)
+        # B : (B, L, N)
+        # C : (B, L, N)
+        # D : (ED)
+
+        # y : (B, L, ED)
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, L, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # (B, L, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, L, ED, N)
+        
+        hs = pscan(deltaA, BX)
+
+        y = (hs @ C.unsqueeze(-1)).squeeze(3) # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)
+
+        y = y + D * x
+
+        return y
+    
+    def selective_scan_seq(self, x, delta, A, B, C, D):
+        # x : (B, L, ED)
+        # Δ : (B, L, ED)
+        # A : (ED, N)
+        # B : (B, L, N)
+        # C : (B, L, N)
+        # D : (ED)
+
+        # y : (B, L, ED)
+
+        _, L, _ = x.shape
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, L, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # (B, L, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, L, ED, N)
+
+        h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) # (B, ED, N)
+        hs = []
+
+        for t in range(0, L):
+            h = deltaA[:, t] * h + BX[:, t]
+            hs.append(h)
+            
+        hs = torch.stack(hs, dim=1) # (B, L, ED, N)
+
+        y = (hs @ C.unsqueeze(-1)).squeeze(3) # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)
+
+        y = y + D * x
+
+        return y
+    
+    # -------------------------- inference -------------------------- #
+    """
+    Concerning auto-regressive inference
+
+    The cool part of using Mamba : inference is constant wrt to sequence length
+    We just have to keep in cache, for each layer, two things :
+    - the hidden state h (which is (B, ED, N)), as you typically would when doing inference with a RNN
+    - the last d_conv-1 inputs of the layer, to be able to compute the 1D conv which is a convolution over the time dimension
+      (d_conv is fixed so this doesn't incur a growing cache as we progress on generating the sequence)
+      (and d_conv is usually very small, like 4, so we just have to "remember" the last 3 inputs)
+
+    Concretely, these two quantities are put inside a cache tuple, and are named h and inputs respectively.
+    h is (B, ED, N), and inputs is (B, ED, d_conv-1)
+    The MambaBlock.step() receives this cache, and, along with outputing the output, alos outputs the updated cache for the next call.
+
+    The cache object is initialized as follows : (None, torch.zeros()).
+    When h is None, the selective scan function detects it and start with h=0.
+    The torch.zeros() isn't a problem (it's same as just feeding the input, because the conv1d is padded)
+
+    As we need one such cache variable per layer, we store a caches object, which is simply a list of cache object. (See mamba_lm.py)
+    """
+    
+    def step(self, x, cache):
+        # x : (B, D)
+        # cache : (h, inputs)
+                # h : (B, ED, N)
+                # inputs : (B, ED, d_conv-1)
+        
+        # y : (B, D)
+        # cache : (h, inputs)
+        
+        h, inputs = cache
+        
+        xz = self.in_proj(x) # (B, 2*ED)
+        x, z = xz.chunk(2, dim=1) # (B, ED), (B, ED)
+
+        # x branch
+        x_cache = x.unsqueeze(2)
+        x = self.conv1d(torch.cat([inputs, x_cache], dim=2))[:, :, self.config.d_conv-1] # (B, ED)
+
+        x = F.silu(x)
+        y, h = self.ssm_step(x, h)
+
+        # z branch
+        z = F.silu(z)
+
+        output = y * z
+        output = self.out_proj(output) # (B, D)
+
+        # prepare cache for next call
+        inputs = torch.cat([inputs[:, :, 1:], x_cache], dim=2) # (B, ED, d_conv-1)
+        cache = (h, inputs)
+        
+        return output, cache
+
+    def ssm_step(self, x, h):
+        # x : (B, ED)
+        # h : (B, ED, N)
+
+        # y : (B, ED)
+        # h : (B, ED, N)
+
+        A = -torch.exp(self.A_log.float()) # (ED, N) # todo : ne pas le faire tout le temps, puisque c'est indépendant de la timestep
+        D = self.D.float()
+        # TODO remove .float()
+
+        deltaBC = self.x_proj(x) # (B, dt_rank+2*N)
+
+        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, dt_rank), (B, N), (B, N)
+        delta = F.softplus(self.dt_proj(delta)) # (B, ED)
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(1) # (B, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, ED, N)
+
+        if h is None:
+            h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) # (B, ED, N)
+
+        h = deltaA * h + BX # (B, ED, N)
+
+        y = (h @ C.unsqueeze(-1)).squeeze(2) # (B, ED, N) @ (B, N, 1) -> (B, ED, 1)
+
+        y = y + D * x
+
+        # todo : pq h.squeeze(1) ??
+        return y, h.squeeze(1)
+
+# taken straight from https://github.com/johnma2006/mamba-minimal/blob/master/model.py
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-5):
+        super().__init__()
+
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x):
+        output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+        return output

models/pscan.py

@@ -0,0 +1,232 @@
+"""
+Mamba: Linear-Time Sequence Modeling with Selective State Spaces
+Copyright (c) Carnegie Mellon University.
+Implemented by alxndrTL from https://github.com/alxndrTL/mamba.py
+"""
+
+
+import math
+import torch
+import torch.nn.functional as F
+
+"""
+
+An implementation of the parallel scan operation in PyTorch (Blelloch version).
+Please see docs/pscan.ipynb for a detailed explanation of what happens here.
+
+"""
+
+def npo2(len):
+    """
+    Returns the next power of 2 above len
+    """
+
+    return 2 ** math.ceil(math.log2(len))
+
+def pad_npo2(X):
+    """
+    Pads input length dim to the next power of 2
+
+    Args:
+        X : (B, L, D, N)
+
+    Returns:
+        Y : (B, npo2(L), D, N)
+    """
+
+    len_npo2 = npo2(X.size(1))
+    pad_tuple = (0, 0, 0, 0, 0, len_npo2 - X.size(1))
+    return F.pad(X, pad_tuple, "constant", 0)
+
+class PScan(torch.autograd.Function):
+    @staticmethod
+    def pscan(A, X):
+        # A : (B, D, L, N)
+        # X : (B, D, L, N)
+
+        # modifies X in place by doing a parallel scan.
+        # more formally, X will be populated by these values :
+        # H[t] = A[t] * H[t-1] + X[t] with H[0] = 0
+        # which are computed in parallel (2*log2(T) sequential steps (ideally), instead of T sequential steps)
+
+        # only supports L that is a power of two (mainly for a clearer code)
+        
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # up sweep (last 2 steps unfolded)
+        Aa = A
+        Xa = X
+        for _ in range(num_steps-2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+            
+            Xa[:, :, :, 1].add_(Aa[:, :, :, 1].mul(Xa[:, :, :, 0]))
+            Aa[:, :, :, 1].mul_(Aa[:, :, :, 0])
+
+            Aa = Aa[:, :, :, 1]
+            Xa = Xa[:, :, :, 1]
+
+        # we have only 4, 2 or 1 nodes left
+        if Xa.size(2) == 4:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            Aa[:, :, 1].mul_(Aa[:, :, 0])
+
+            Xa[:, :, 3].add_(Aa[:, :, 3].mul(Xa[:, :, 2] + Aa[:, :, 2].mul(Xa[:, :, 1])))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            return
+        else:
+            return
+
+        # down sweep (first 2 steps unfolded)
+        Aa = A[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)]
+        Xa = X[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)]
+        Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 1]))
+        Aa[:, :, 2].mul_(Aa[:, :, 1])
+
+        for k in range(num_steps-3, -1, -1):
+            Aa = A[:, :, 2**k-1:L:2**k]
+            Xa = X[:, :, 2**k-1:L:2**k]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+
+            Xa[:, :, 1:, 0].add_(Aa[:, :, 1:, 0].mul(Xa[:, :, :-1, 1]))
+            Aa[:, :, 1:, 0].mul_(Aa[:, :, :-1, 1])
+
+    @staticmethod
+    def pscan_rev(A, X):
+        # A : (B, D, L, N)
+        # X : (B, D, L, N)
+
+        # the same function as above, but in reverse
+        # (if you flip the input, call pscan, then flip the output, you get what this function outputs)
+        # it is used in the backward pass
+
+        # only supports L that is a power of two (mainly for a clearer code)
+
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # up sweep (last 2 steps unfolded)
+        Aa = A
+        Xa = X
+        for _ in range(num_steps-2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+                    
+            Xa[:, :, :, 0].add_(Aa[:, :, :, 0].mul(Xa[:, :, :, 1]))
+            Aa[:, :, :, 0].mul_(Aa[:, :, :, 1])
+
+            Aa = Aa[:, :, :, 0]
+            Xa = Xa[:, :, :, 0]
+
+        # we have only 4, 2 or 1 nodes left
+        if Xa.size(2) == 4:
+            Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 3]))
+            Aa[:, :, 2].mul_(Aa[:, :, 3])
+
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1].add(Aa[:, :, 1].mul(Xa[:, :, 2]))))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1]))
+            return
+        else:
+            return
+
+        # down sweep (first 2 steps unfolded)
+        Aa = A[:, :, 0:L:2**(num_steps-2)]
+        Xa = X[:, :, 0:L:2**(num_steps-2)]
+        Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 2]))
+        Aa[:, :, 1].mul_(Aa[:, :, 2])
+
+        for k in range(num_steps-3, -1, -1):
+            Aa = A[:, :, 0:L:2**k]
+            Xa = X[:, :, 0:L:2**k]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+
+            Xa[:, :, :-1, 1].add_(Aa[:, :, :-1, 1].mul(Xa[:, :, 1:, 0]))
+            Aa[:, :, :-1, 1].mul_(Aa[:, :, 1:, 0])
+
+    @staticmethod
+    def forward(ctx, A_in, X_in):
+        """
+        Applies the parallel scan operation, as defined above. Returns a new tensor.
+        If you can, privilege sequence lengths that are powers of two.
+
+        Args:
+            A_in : (B, L, D, N)
+            X_in : (B, L, D, N)
+
+        Returns:
+            H : (B, L, D, N)
+        """
+
+        L = X_in.size(1)
+
+        # cloning is requiered because of the in-place ops
+        if L == npo2(L):
+            A = A_in.clone()
+            X = X_in.clone()
+        else:
+            # pad tensors (and clone btw)
+            A = pad_npo2(A_in) # (B, npo2(L), D, N)
+            X = pad_npo2(X_in) # (B, npo2(L), D, N)
+        
+        # prepare tensors
+        A = A.transpose(2, 1) # (B, D, npo2(L), N)
+        X = X.transpose(2, 1) # (B, D, npo2(L), N)
+
+        # parallel scan (modifies X in-place)
+        PScan.pscan(A, X)
+
+        ctx.save_for_backward(A_in, X)
+        
+        # slice [:, :L] (cut if there was padding)
+        return X.transpose(2, 1)[:, :L]
+    
+    @staticmethod
+    def backward(ctx, grad_output_in):
+        """
+        Flows the gradient from the output to the input. Returns two new tensors.
+
+        Args:
+            ctx : A_in : (B, L, D, N), X : (B, D, L, N)
+            grad_output_in : (B, L, D, N)
+
+        Returns:
+            gradA : (B, L, D, N), gradX : (B, L, D, N)
+        """
+
+        A_in, X = ctx.saved_tensors
+
+        L = grad_output_in.size(1)
+
+        # cloning is requiered because of the in-place ops
+        if L == npo2(L):
+            grad_output = grad_output_in.clone()
+            # the next padding will clone A_in
+        else:
+            grad_output = pad_npo2(grad_output_in) # (B, npo2(L), D, N)
+            A_in = pad_npo2(A_in) # (B, npo2(L), D, N)
+
+        # prepare tensors
+        grad_output = grad_output.transpose(2, 1)
+        A_in = A_in.transpose(2, 1) # (B, D, npo2(L), N)
+        A = torch.nn.functional.pad(A_in[:, :, 1:], (0, 0, 0, 1)) # (B, D, npo2(L), N) shift 1 to the left (see hand derivation)
+
+        # reverse parallel scan (modifies grad_output in-place)
+        PScan.pscan_rev(A, grad_output)
+
+        Q = torch.zeros_like(X)
+        Q[:, :, 1:].add_(X[:, :, :-1] * grad_output[:, :, 1:])
+
+        return Q.transpose(2, 1)[:, :L], grad_output.transpose(2, 1)[:, :L]
+    
+pscan = PScan.apply

models/time_transformer

@@ -0,0 +1,256 @@
+'''
+Exploring Temporal Coherence for More General Video Face Forgery Detection @ ICCV'2021
+Copyright (c) Xiamen University and its affiliates.
+Modified by Yinglin Zheng from https://github.com/yinglinzheng/FTCN
+'''
+
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + x
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, mask = None):
+        b, n, _, h = *x.shape, self.heads
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
+
+        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+        mask_value = -torch.finfo(dots.dtype).max
+
+        if mask is not None:
+            mask = F.pad(mask.flatten(1), (1, 0), value = True)
+            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
+            mask = rearrange(mask, 'b i -> b () i ()') * rearrange(mask, 'b j -> b () () j')
+            dots.masked_fill_(~mask, mask_value)
+            del mask
+
+        attn = dots.softmax(dim=-1)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        out =  self.to_out(out)
+        return out
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
+                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
+            ]))
+    def forward(self, x, mask = None):
+        for attn, ff in self.layers:
+            x = attn(x, mask = mask)
+            x = ff(x)
+        return x
+
+class ViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
+        super().__init__()
+        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
+        num_patches = (image_size // patch_size) ** 2
+        patch_dim = channels * patch_size ** 2
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
+            nn.Linear(patch_dim, dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.pool = pool
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, img, mask = None):
+        x = self.to_patch_embedding(img)
+        b, n, _ = x.shape #batch,num_patches,channels  #
+
+        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.transformer(x, mask)
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return self.mlp_head(x)
+
+
+
+def valid_idx(idx, h):
+    i = idx // h
+    j = idx % h
+    pad = h // 7
+    if j < pad or i >= h - pad or j >= h - pad:
+        return False
+    else:
+        return True
+
+import random
+from math import sqrt
+class RandomSelect(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        # batch,7x7
+        size=x.shape[1]
+        h=int(sqrt(size))
+        candidates = list(range(size))
+        candidates = [idx for idx in candidates if valid_idx(idx, h)]
+        max_k = len(candidates)
+        if self.training:
+            k = 8
+            if k==-1:
+                k=max_k
+        else:
+            k = max_k
+        candidates = random.sample(candidates, k)
+        x = x[:,candidates]
+        return x
+
+class VideoiT(nn.Module):
+    def __init__(self, *, image_size, patch_size, num_patches, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
+        super().__init__()
+        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
+        patch_dim = channels * patch_size ** 2
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.to_patch = Rearrange('b c t (h p1) (w p2) -> b (h w) t (p1 p2 c)', p1 = patch_size, p2 = patch_size)
+        self.patch_to_embedding=nn.Linear(patch_dim, dim)
+        self.num_patches=num_patches
+        
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.pool = pool
+        self.random_select=RandomSelect()
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, img, mask = None):
+        real_b=img.shape[0]
+        x = self.to_patch(img)
+        x = self.random_select(x)
+        n=x.shape[1]
+        x=x.reshape(real_b*n,self.num_patches,-1)
+        x = self.patch_to_embedding(x)
+        b, n, _ = x.shape #batch,num_patches,channels  #
+
+        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.transformer(x, mask)
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        x = self.mlp_head(x)
+        x = x.reshape(real_b,-1)
+        return x
+
+
+class TimeTransformer(nn.Module):
+    def __init__(self,num_patches, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', dim_head = 64, dropout = 0., emb_dropout = 0.):
+        super().__init__()
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.num_patches=num_patches
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.pool = pool
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, x):
+        b, n, _ = x.shape #batch,num_patches,channels  #
+
+        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.transformer(x, mask=None)
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return self.mlp_head(x)

requirements.txt

@@ -0,0 +1,59 @@
+absl-py==2.3.1
+albucore==0.0.24
+albumentations==2.0.8
+annotated-types==0.7.0
+certifi==2025.10.5
+charset-normalizer==3.4.3
+eval-type-backport==0.2.2
+filelock==3.19.1
+fsspec==2025.9.0
+ftfy==6.3.1
+grpcio==1.75.1
+hf-xet==1.1.10
+huggingface-hub==0.35.3
+idna==3.10
+importlib-metadata==8.7.0
+jinja2==3.1.6
+joblib==1.5.2
+markdown==3.9
+markupsafe==3.0.3
+mpmath==1.3.0
+networkx==3.2.1
+numpy==2.0.2
+opencv-python==4.12.0.88
+opencv-python-headless==4.12.0.88
+packaging==25.0
+pandas==2.3.3
+pillow==11.3.0
+protobuf==6.32.1
+pydantic==2.11.10
+pydantic-core==2.33.2
+python-dateutil==2.9.0.post0
+pytz==2025.2
+pyyaml==6.0.3
+regex==2025.9.18
+requests==2.32.5
+safetensors==0.6.2
+scikit-learn==1.6.1
+scipy==1.13.1
+simsimd==6.5.3
+six==1.17.0
+stringzilla==4.1.0
+sympy==1.14.0
+tensorboard==2.20.0
+tensorboard-data-server==0.7.2
+threadpoolctl==3.6.0
+timm==1.0.20
+tokenizers==0.22.1
+torch==2.8.0
+torchvision==0.23.0
+tqdm==4.67.1
+transformers==4.57.0
+triton==3.4.0
+typing-extensions==4.15.0
+typing-inspection==0.4.2
+tzdata==2025.2
+urllib3==2.5.0
+wcwidth==0.2.14
+werkzeug==3.1.3
+zipp==3.23.0

results.csv

@@ -0,0 +1,6 @@
+file_name,predicted_prob,predicted_label
+kling_kling_6ebbdf57a77769ab5a9cfb682249a9cfb9e389666d1b568bedcaf37c45c521f6_f000.jpg,0.8246555924415588,1
+kling_kling_6ebbdf57a77769ab5a9cfb682249a9cfb9e389666d1b568bedcaf37c45c521f6_f001.jpg,0.8141521215438843,1
+kling_kling_6ebbdf57a77769ab5a9cfb682249a9cfb9e389666d1b568bedcaf37c45c521f6_f002.jpg,0.8013387322425842,1
+kling_kling_6ebbdf57a77769ab5a9cfb682249a9cfb9e389666d1b568bedcaf37c45c521f6_f003.jpg,0.7972325086593628,1
+kling_kling_6ebbdf57a77769ab5a9cfb682249a9cfb9e389666d1b568bedcaf37c45c521f6_f004.jpg,0.8086235523223877,1

script.py

@@ -0,0 +1,22 @@
+import subprocess
+
+scripts = [
+    "extract_frames.py",
+    "eval2.py",
+    "create_submission.py"
+]
+
+def run_script(script):
+    print(f"\n🚀 Running {script}...")
+    result = subprocess.run(["python", script], capture_output=True, text=True)
+    if result.returncode == 0:
+        print(f"✅ {script} completed successfully.\n")
+        print(result.stdout)
+    else:
+        print(f"❌ {script} failed with error:\n{result.stderr}")
+        exit(1)
+
+if __name__ == "__main__":
+    for script in scripts:
+        run_script(script)
+    print("🎉 All scripts completed successfully!")

submission.csv

@@ -0,0 +1,2 @@
+id,pred,score
+kling_kling_6ebbdf57a77769ab5a9cfb682249a9cfb9e389666d1b568bedcaf37c45c521f6,generated,0.8092005014419555