init space

.gitignore

@@ -0,0 +1,16 @@
+wandb
+__pycache__
+.vscode
+notebooks
+results
+*.ipynb 
+*.ipynb_checkpoints
+eval_results
+tests
+downloads
+ckpts
+demo_images
+eval/OneIG-Benchmark/models
+eval/OneIG-Benchmark/scripts/style/models
+eval/OneIG-Benchmark/results*
+models/

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -0,0 +1,231 @@
+import logging
+import os
+import random
+import re
+import sys
+import warnings
+import gradio as gr
+import math
+import torch
+import subprocess
+subprocess.run(
+    "pip install flash-attn==2.8.2 --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+try:
+    from modeling.t2i_pipeline import BitDanceT2IPipeline
+except ImportError:
+    print("Warning: Could not import BitDanceT2IPipeline. Please ensure 'modeling' folder is present.")
+
+from huggingface_hub import snapshot_download
+
+save_dir = "models/BitDance-14B-64x"
+repo_id = "shallowdream204/BitDance-14B-64x"
+cache_dir = save_dir + "/cache"
+
+snapshot_download(cache_dir=cache_dir,
+  local_dir=save_dir,
+  repo_id=repo_id,
+  local_dir_use_symlinks=False,
+  resume_download=True,
+  allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
+)
+
+# ==================== Environment Variables ==================================
+MODEL_PATH = save_dir
+
+# =============================================================================
+warnings.filterwarnings("ignore")
+logging.getLogger("transformers").setLevel(logging.ERROR)
+
+# ==================== Resolution Settings ====================================
+RAW_RESOLUTIONS = [
+    [2048, 512],
+    [1920, 512],
+    [1536, 640],
+    [1280, 768],
+    [1152, 896],
+    [1024, 1024],
+    [896, 1152],
+    [768, 1280],
+    [640, 1536],
+    [512, 1920],
+    [512, 2048],
+]
+
+RESOLUTION_CHOICES = []
+for w, h in RAW_RESOLUTIONS:
+    divisor = math.gcd(w, h)
+    ratio_w = w // divisor
+    ratio_h = h // divisor
+    label = f"{w}x{h} ({ratio_w}:{ratio_h})"
+    RESOLUTION_CHOICES.append(label)
+
+DEFAULT_RES = "1024x1024 (1:1)"
+
+EXAMPLE_PROMPTS = [
+    ["一位穿着粉色吊带罗纹长裙的亚洲少女，外搭一件米白色毛绒短开襟衫，在阳光洒落的森林小径上侧身回眸。她拥有淡粉色薰衣草发色的甜美脸庞，发间别着一朵白色小花。黄金时段的光线穿过浓密的树叶，在深绿色的背景上形成美丽的景深光斑 和柔和光晕。电影级肖像摄影，超高画质，细腻的皮肤纹理，强调少女的温柔与唯美浪漫的日系氛围。"],
+    ["一幅具有电影感的胶片肖像，一位美丽的中国女生，凌乱的黑发在风中飘动遮住脸庞，眼神灵动地看着镜头。她在画面的左1/3处。她围着一条厚实的鲜红色针织围巾，穿着一件破旧的米色羊羔毛外套。背景是日落时分寒冷、干枯的荒野和远山。强烈的金色逆光直射镜头，产生巨大的镜头眩光和朦胧的光晕效果，空气中有尘埃感。胶片颗粒质感，浅景深，自然原始的风格。"],
+    ["一个半人半机械的黑客，坐在充满全息屏幕的黑暗房间里，绿色的代码光映照在他的脸上，赛博朋克风格，高科技细节，锐利的焦点。"],
+    ["A surreal double exposure portrait that blends a woman’s face with a beautiful seascape. The overall mood is dreamy and mystical, with rich colors and intricate details."],
+    ["A close-up, macro photography stock photo of a strawberry intricately sculpted into the shape of a hummingbird in mid-flight, its wings a blur as it sips nectar from a vibrant, tubular flower. The backdrop features a lush, colorful garden with a soft, bokeh effect, creating a dreamlike atmosphere. The image is exceptionally detailed and captured with a shallow depth of field, ensuring a razor-sharp focus on the strawberry-hummingbird and gentle fading of the background. The high resolution, professional photographers style, and soft lighting illuminate the scene in a very detailed manner, professional color grading amplifies the vibrant colors and creates an image with exceptional clarity. The depth of field makes the hummingbird and flower stand out starkly against the bokeh background."],
+    ["网红咖啡店内部，透过钢化玻璃拍摄，中景平视角度；玻璃表面有环境反光与色彩叠影，人物面部柔光打亮，坐着看向镜头，穿着带大毛领的宽松上衣；白天咖啡店，太阳光线打在人物脸上，玻璃反光清透自然，ccd质感。"],
+    ["室内中景人像摄影，复古胶片风格，电影叙事感画面。一位清纯气质的年轻女性，留着黑色齐刘海长直发，妆容清透伪素颜，皮肤白皙透亮。她身穿一件质地柔软、淡绿色的马海毛（Mohair）绒毛毛衣，质感毛绒蓬松，下身搭配淡青色棉麻长裙。人物慵懒地蜷缩/侧卧在沙发角落，身体姿态放松柔软，呈现自然的C型曲线。一只手轻轻拿着一颗鲜红的番茄靠近脸颊和下巴，眼神迷离、温柔且深情地直视镜头，表情处于放空与凝视之间，极具故事感。复古文艺的室内一角，沙发上铺着淡雅的复古碎花布艺沙发罩，身旁放着一盘红色的番茄作为前景点缀。背景虚化，隐约可见室内的陈设与绿植，整体环境色调偏向青绿色���胶片感。极具艺术感的局部自然光（丁达尔效应光斑）。一束明亮的午后阳光精准地照射在手部、手中的番茄以及面部一侧，形成强烈的明暗对比（Chiaroscuro）。高光部分带有光晕（Bloom），阴影部分呈现胶片特有的青蓝色调，光影层次丰富。。慵懒、静谧、梦幻、日系文艺、情绪感强、高级且富有夏末秋初的诗意。。模拟胶片相机（如Contax T3或Pentax 67）拍摄，使用50mm标准定焦镜头，大光圈（f/1.8）制造柔和的背景虚化。后期加入明显的粗颗粒胶片滤镜（Heavy Film Grain）和色彩偏移，增强模拟摄影的真实感与年代感。极度真实的皮肤质感，保留面部微小的毛孔和纹理，拒绝过度磨皮；马海毛毛衣在逆光下呈现出清晰的绒毛光晕边缘；番茄表面光滑的高光反射；碎花布料的褶皱细节；整体画面覆盖一层复古的胶片噪点。"],
+]
+
+def get_resolution(resolution_str):
+    match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution_str)
+    if match:
+        return int(match.group(1)), int(match.group(2))
+    return 1024, 1024
+
+def load_models(model_path):
+    print(f"Loading BitDance model from {model_path}...")
+    
+    if not os.path.exists(model_path):
+        print(f"Warning: Model path {model_path} does not exist locally. Attempting to load anyway (or handle download logic here).")
+
+    pipe = BitDanceT2IPipeline(model_path=model_path, device="cuda")
+    return pipe
+
+def generate_image(
+    pipe,
+    prompt,
+    resolution,
+    seed=42,
+    guidance_scale=7.5,
+    num_inference_steps=50,
+):
+    width, height = get_resolution(resolution)
+    
+    images = pipe.generate(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_sampling_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_images=1,
+        seed=seed
+    )
+    
+    return images[0]
+
+pipe = None
+
+def init_app():
+    global pipe
+    try:
+        pipe = load_models(MODEL_PATH)
+        print("Model loaded successfully.")
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        pipe = None
+
+def generate(
+    prompt,
+    resolution,
+    seed=42,
+    steps=50,
+    guidance_scale=7.5,
+    random_seed=True,
+    gallery_images=None,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if random_seed:
+        new_seed = random.randint(1, 1000000)
+    else:
+        new_seed = seed if seed != -1 else random.randint(1, 1000000)
+
+    if pipe is None:
+        raise gr.Error("Model not loaded.")
+
+    print(f"Generating: Prompt='{prompt[:20]}...', Res={resolution}, Seed={new_seed}, Steps={steps}, CFG={guidance_scale}")
+
+    try:
+        image = generate_image(
+            pipe=pipe,
+            prompt=prompt,
+            resolution=resolution,
+            seed=new_seed,
+            guidance_scale=guidance_scale,
+            num_inference_steps=int(steps),
+        )
+    except Exception as e:
+        raise gr.Error(f"Generation failed: {str(e)}")
+
+    if gallery_images is None:
+        gallery_images = []
+    
+    gallery_images = [image] + gallery_images 
+    
+    return gallery_images, str(new_seed), int(new_seed)
+
+init_app()
+
+# ==================== Gradio UI ====================
+
+with gr.Blocks(title="BitDance Demo") as demo:
+    gr.Markdown(
+        """<div align="center">
+
+# BitDance Generation Demo
+
+[![GitHub](https://img.shields.io/badge/GitHub-BitDance-181717?logo=github&logoColor=white)](https://github.com/shallowdream204/BitDance)
+
+*BitDance: Scaling Autoregressive Generative Models with Binary Tokens*
+
+</div>"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt here...")
+            
+            resolution = gr.Dropdown(
+                value=DEFAULT_RES, 
+                choices=RESOLUTION_CHOICES, 
+                label="Resolution (Width x Height)"
+            )
+
+            with gr.Row():
+                seed = gr.Number(label="Seed", value=42, precision=0)
+                random_seed = gr.Checkbox(label="Random Seed", value=True)
+
+            with gr.Row():
+                steps = gr.Slider(label="Diffusion Sampling Steps", minimum=10, maximum=100, value=50, step=1)
+                guidance_scale = gr.Slider(label="CFG Guidance Scale", minimum=1.0, maximum=15.0, value=7.5, step=0.5)
+
+            generate_btn = gr.Button("Generate", variant="primary")
+
+            gr.Markdown("### 📝 Example Prompts")
+            gr.Examples(examples=EXAMPLE_PROMPTS, inputs=prompt_input, label=None)
+
+        with gr.Column(scale=1):
+            output_gallery = gr.Gallery(
+                label="Generated Images",
+                columns=2,
+                rows=2,
+                height=600,
+                object_fit="contain",
+                format="png",
+                interactive=False,
+            )
+            used_seed = gr.Textbox(label="Seed Used", interactive=False)
+
+    generate_btn.click(
+        generate,
+        inputs=[prompt_input, resolution, seed, steps, guidance_scale, random_seed, output_gallery],
+        outputs=[output_gallery, used_seed, seed],
+        api_visibility="public",
+    )
+
+css = """
+.fillable{max-width: 1230px !important}
+"""
+
+if __name__ == "__main__":
+    demo.launch(css=css, mcp_server=True)

modeling/t2i_pipeline.py

@@ -0,0 +1,283 @@
+import torch
+from torch import nn
+from einops import rearrange
+
+from transformers import set_seed
+from PIL import Image
+
+import numpy as np
+from torch import nn
+from transformers import AutoTokenizer, Qwen3ForCausalLM, Qwen3Config
+
+from modeling.utils import MLPconnector
+from modeling.vision_encoder.autoencoder import VQModel
+from modeling.vision_head.flow_head_parallel_x import DiffHead
+
+from safetensors.torch import load_file as load_sft
+import json
+import os
+from tqdm import tqdm
+
+IMAGE_SIZE_LIST = [
+    # --- 1024px Area ---
+    [2048, 512],
+    [1920, 512],
+    [1536, 640],
+    [1280, 768],
+    [1152, 896],
+    [1024, 1024],
+    [896, 1152],
+    [768, 1280],
+    [640, 1536],
+    [512, 1920],
+    [512, 2048],
+    # --- 512px Area ---
+    [1024, 256],
+    [896, 256],
+    [640, 384],
+    [512, 512],
+    [384, 640],
+    [256, 896],
+    [256, 1024],
+]
+
+class BitDanceT2IPipeline:
+    def __init__(self, model_path, device='cuda'):
+        self.device = device
+        # LLM and tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.llm_config = Qwen3Config.from_pretrained(model_path)
+        self.llm_model = Qwen3ForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).eval().to(device)
+        self.hidden_size = self.llm_config.hidden_size
+
+        # Autoencoder
+        with open(os.path.join(model_path, 'ae_config.json'), "r") as f:
+            self.ae_config = json.load(f)
+        self.ae = VQModel(**self.ae_config).eval()
+        self.ae.load_state_dict(load_sft(os.path.join(model_path, 'ae.safetensors')), strict=True, assign=True)
+        self.ae.to(device)
+        self.vae_patch_size = 2 ** (len(self.ae_config['ddconfig']['ch_mult'])-1)
+
+        # Vision head
+        with open(os.path.join(model_path, 'vision_head_config.json'), "r") as f:
+            self.vision_head_config = json.load(f)
+        self.vision_head = DiffHead(**self.vision_head_config).eval()
+        self.vision_head.load_state_dict(load_sft(os.path.join(model_path, 'vision_head.safetensors')), strict=True, assign=True)
+        self.vision_head.to(device)
+        self.parallel_num = self.vision_head_config['parallel_num']
+        print(f'use {self.parallel_num}-token parallel prediction per step')
+        self.ps = int(self.parallel_num ** 0.5)
+
+        # Projector
+        self.embed_vision_mlp = MLPconnector(self.ae_config['ddconfig']['z_channels'], self.hidden_size, "gelu_pytorch_tanh")
+        self.embed_vision_mlp.load_state_dict(load_sft(os.path.join(model_path, 'projector.safetensors')), strict=True, assign=True)
+        self.embed_vision_mlp.to(device)
+
+        # 2D sinusoidal position embedding
+        self.build_pos_embed()
+
+    def build_pos_embed(self, max_len=4096):
+        max_len = max_len // self.vae_patch_size
+        pos_embed_1d = self._get_1d_sincos_pos_embed(self.hidden_size//2, max_len)
+        pos_embed_1d = nn.Parameter(pos_embed_1d, requires_grad=False)
+        self.pos_embed_1d = pos_embed_1d.to(self.device)
+
+    def _get_1d_sincos_pos_embed(self, dim, max_len, pe_interpolation=1.0):
+        assert dim % 2 == 0
+        omega = torch.arange(dim // 2, dtype=torch.float32)
+        omega /= dim / 2.0
+        omega = 1.0 / 10000**omega  # (D/4,)
+
+        pos = torch.arange(max_len, dtype=torch.float32) / pe_interpolation
+        out = torch.einsum("m,d->md", pos, omega)  # (max_len, D/4)
+
+        emb_sin = torch.sin(out)
+        emb_cos = torch.cos(out)
+        return torch.cat([emb_sin, emb_cos], dim=1)  # (max_len, D/2)
+
+    def get_2d_embed(self, h, w, ps=1):
+        emb_v = self.pos_embed_1d[:h]
+        emb_h = self.pos_embed_1d[:w]
+
+        grid_v = emb_v.view(h, 1, self.hidden_size//2).repeat(1, w, 1)
+        grid_h = emb_h.view(1, w, self.hidden_size//2).repeat(h, 1, 1)
+
+        pos_embed = torch.cat([grid_h, grid_v], dim=-1) # h w c
+
+        return rearrange(pos_embed, '(h p1) (w p2) c -> (h w p1 p2) c', p1=ps, p2=ps)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompt: str,
+        height: int = 1024,
+        width: int = 1024,
+        num_sampling_steps: int = 50,
+        guidance_scale: float = 7.5,
+        num_images: int = 1,
+        seed: int = 1234,
+    ):
+        # Set seed for reproducibility
+        if seed is not None:
+            set_seed(seed)
+        # Calculate max_length dynamically based on image_size and stride of 16
+        max_length = (height // self.vae_patch_size) * (width // self.vae_patch_size)
+        
+        image_size = [height, width]
+        if image_size not in IMAGE_SIZE_LIST:
+            raise ValueError(f"image_size {image_size} is not supported. Please choose from {IMAGE_SIZE_LIST}")
+
+        with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+            gen_images = self.gen_image(
+                cond_prompt=f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n",
+                uncond_prompt="<|im_start|>assistant\n",
+                guidance_scale=guidance_scale,
+                num_sampling_steps=num_sampling_steps,
+                num_images=num_images,
+                image_size=image_size,
+                max_length=max_length,
+                show_progress=True,
+            )
+        
+        gen_images = (
+            torch.clamp(127.5 * gen_images + 128.0, 0, 255)
+            .permute(0, 2, 3, 1)
+            .to("cpu", dtype=torch.uint8)
+            .numpy()
+        )
+        pil_images = []
+        for i in range(gen_images.shape[0]):
+            img_array = gen_images[i]
+            if img_array.dtype != np.uint8:
+                img_array = img_array.astype(np.uint8)
+            pil_images.append(Image.fromarray(img_array))
+
+        return pil_images
+
+    @torch.no_grad()
+    def gen_image(self,
+        cond_prompt,
+        uncond_prompt=None,
+        guidance_scale: float = 1.0,
+        num_sampling_steps: int = 50,
+        max_length: int = 64,
+        num_images: int = 1,
+        image_size = [256, 256],
+        show_progress: bool = False,
+    ):
+        tokenizer = self.tokenizer
+        device = self.device
+        model = self.llm_model.model
+
+        step_width = self.parallel_num
+        num_steps = max_length // step_width
+
+        cond_ids = torch.tensor(tokenizer.encode(cond_prompt), device=device, dtype=torch.long)
+        cond_emb = model.embed_tokens(cond_ids)
+        if guidance_scale > 1.0:
+            uncond_ids = torch.tensor(tokenizer.encode(uncond_prompt), device=device, dtype=torch.long)
+            uncond_emb = model.embed_tokens(uncond_ids)
+
+        img_start_id = tokenizer.convert_tokens_to_ids("<|vision_start|>")
+        res_h_token_id = tokenizer.convert_tokens_to_ids(f"<|res_{image_size[0] // self.vae_patch_size}|>")
+        res_w_token_id = tokenizer.convert_tokens_to_ids(f"<|res_{image_size[1] // self.vae_patch_size}|>")
+        img_start_emb = model.embed_tokens(torch.tensor([img_start_id, res_h_token_id, res_w_token_id], device=device))
+
+        h, w = image_size[0] // self.vae_patch_size, image_size[1] // self.vae_patch_size
+        # prepare diff pos embed
+        pos_embed_for_diff = self.get_2d_embed(h, w, ps=self.ps if hasattr(self, 'ps') else 1).unsqueeze(0)
+
+        # add query tokens for parallel decoding
+        for i in range(1, self.parallel_num):
+            query_token = torch.tensor([tokenizer.convert_tokens_to_ids(f"<|query_{i}|>")], device=self.device, dtype=torch.long)
+            query_embed = self.llm_model.model.embed_tokens(query_token)
+            img_start_emb = torch.cat([img_start_emb, query_embed], dim=0)
+
+        input_embeds_cond = torch.cat(
+            [cond_emb, img_start_emb], dim=0
+        ).unsqueeze(0).repeat(num_images, 1, 1)
+        outputs_c = model(
+            inputs_embeds=input_embeds_cond[:, :-step_width, :],
+            use_cache=True,
+        )
+        pkv_c = outputs_c.past_key_values
+
+        # bidirectional attn
+        bi_attn_mask = torch.ones(
+                (input_embeds_cond.shape[0], 1, step_width, step_width+pkv_c[0][0].shape[2]),
+                dtype=torch.bool,
+                device=device,
+            )
+        outputs_c = model(
+            inputs_embeds=input_embeds_cond[:, -step_width:, :],
+            past_key_values=pkv_c,
+            use_cache=True,
+            attention_mask=bi_attn_mask,
+        )
+        pkv_c = outputs_c.past_key_values
+        hidden_c = outputs_c.last_hidden_state[:, -step_width:] # [B, parallel_num, D]
+
+        if guidance_scale > 1.0:
+            input_embeds_uncond = torch.cat(
+                [uncond_emb, img_start_emb], dim=0
+            ).unsqueeze(0).repeat(num_images, 1, 1)
+            outputs_u = model(
+                inputs_embeds=input_embeds_uncond[:, :-step_width, :],
+                use_cache=True,
+            )
+            pkv_u = outputs_u.past_key_values
+            outputs_u = model(
+                inputs_embeds=input_embeds_uncond[:, -step_width:, :],
+                past_key_values=pkv_u,
+                use_cache=True,
+                attention_mask=bi_attn_mask,
+            )
+            pkv_u = outputs_u.past_key_values
+            hidden_u = outputs_u.last_hidden_state[:, -step_width:] # [B, parallel_num, D]
+
+        out_tokens = []
+        if show_progress:
+            pbar = tqdm(total=num_steps, desc="Decoding Steps")
+        for step in range(num_steps):
+            if show_progress:
+                pbar.update(1)
+            h_fused = torch.cat([hidden_c, hidden_u], dim=0) if guidance_scale > 1.0 else hidden_c
+            h_fused = h_fused + pos_embed_for_diff[:, step*step_width:(step+1)*step_width, :]
+            pred_latents = self.vision_head.sample(h_fused, num_sampling_steps=num_sampling_steps, cfg=guidance_scale)
+            # important! LFQ is used here
+            curr_tokens = torch.sign(pred_latents)
+            curr_embeds = self.embed_vision_mlp(curr_tokens)
+            out_tokens.append(curr_tokens[:num_images])
+            model_input = curr_embeds # [B, N, D]
+            # 2d pos embed
+            model_input = model_input + pos_embed_for_diff[:, step*step_width:(step+1)*step_width, :]
+ 
+            # bidirectional attn mask
+            bi_attn_mask = torch.ones(
+                (model_input.shape[0], 1, model_input.shape[1], model_input.shape[1]+pkv_c[0][0].shape[2]), 
+                dtype=torch.bool,
+                device=device
+            )
+            outputs_c = model(inputs_embeds=model_input[:num_images], past_key_values=pkv_c, use_cache=True, attention_mask=bi_attn_mask[:num_images])
+            pkv_c = outputs_c.past_key_values
+            hidden_c = outputs_c.last_hidden_state[:, -step_width:] # [B, parallel_num, D]
+
+            if guidance_scale > 1.0:
+                outputs_u = model(inputs_embeds=model_input[num_images:], past_key_values=pkv_u, use_cache=True, attention_mask=bi_attn_mask[num_images:])
+                pkv_u = outputs_u.past_key_values
+                hidden_u = outputs_u.last_hidden_state[:, -step_width:] # [B, parallel_num, D]
+
+        full_output = torch.cat(out_tokens, dim=1)
+        image = self.decode_image(full_output, [h, w], ps=self.ps if hasattr(self, 'ps') else 1) # [num_images, c, h, w]
+        return image
+
+    def decode_image(self, image_latents, image_size=None, ps=1):
+        if image_size is None:
+            h = w = int(image_latents.size(1) ** 0.5)
+        else:
+            h, w = image_size
+            
+        image_latents = rearrange(image_latents, 'b (h w p1 p2) c -> b c (h p1) (w p2)', h=h//ps, w=w//ps, p1=ps, p2=ps)
+        output = self.ae.decode(image_latents)     # [1, c, h, w]
+        
+        return output

modeling/utils.py

@@ -0,0 +1,216 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from torch.nn.attention.flex_attention import or_masks, and_masks
+
+from transformers.activations import ACT2FN
+
+class MLPconnector(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, hidden_act: str):
+        super().__init__()
+        self.activation_fn = ACT2FN[hidden_act]
+        self.fc1 = nn.Linear(in_dim, out_dim)
+        self.fc2 = nn.Linear(out_dim, out_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+def create_sparse_mask(document_lens, split_lens, attn_modes, parallel_num, device):
+    parallel_causal_num = 2
+    parallel_block_causal_num = parallel_num
+
+    def causal_mask(b, h, q_idx, kv_idx):
+        return q_idx >= kv_idx
+
+    def parallel_block_mask(b, h, q_idx, kv_idx):
+        same_seg = segment_ids[q_idx] == segment_ids[kv_idx]
+        is_par = is_parallel[q_idx] 
+        
+        lq = local_ids[q_idx]
+        lk = local_ids[kv_idx]
+        
+        in_block_region = (lq >= parallel_causal_num) & (lk >= parallel_causal_num)
+        
+        same_block = ((lq - parallel_causal_num) // parallel_block_causal_num) == ((lk - parallel_causal_num) // parallel_block_causal_num)
+        
+        return same_seg & is_par & in_block_region & same_block
+
+    def sample_mask(b, h, q_idx, kv_idx):
+        return document_id[q_idx] == document_id[kv_idx]
+
+    segment_ids_list = []
+    local_ids_list = []
+    is_parallel_list = []
+    
+    current_seg_id = 0
+    for length, mode in zip(split_lens, attn_modes):
+        segment_ids_list.extend([current_seg_id] * length)
+        local_ids_list.extend(list(range(length)))
+        is_parallel_list.extend([True if mode == 'parallel' else False] * length)
+        current_seg_id += 1
+
+    segment_ids = torch.tensor(segment_ids_list, device=device, dtype=torch.long)
+    local_ids = torch.tensor(local_ids_list, device=device, dtype=torch.long)
+    is_parallel = torch.tensor(is_parallel_list, device=device, dtype=torch.bool)
+
+    document_id = torch.cat([torch.full((l,), i, device=device) for i, l in enumerate(document_lens, start=1)])
+
+    return and_masks(or_masks(causal_mask, parallel_block_mask), sample_mask)
+
+def top_k_top_p_filtering(
+    logits,
+    top_k: int = 0,
+    top_p: float = 1.0,
+    filter_value: float = -float("Inf"),
+    min_tokens_to_keep: int = 1,
+):
+    """Filter a distribution of logits using top-k and/or top-p (nucleus) filtering."""
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+
+    return logits
+
+
+def sample_codebook(
+    pred_logits,
+    cur_item_type,
+    codebook,
+    do_sample: bool = True,
+    temperature: float = 1.0,
+    top_k: int = 0,
+    top_p: float = 1.0,
+):
+    """
+    pred_logits: (B, vocab_size)
+    cur_item_type: 'text' or 'vision'
+    """
+    # 1. Apply temperature
+    logits = pred_logits / max(temperature, 1e-5)
+
+    # 2. Apply top-k / top-p filtering
+    if top_k > 0 or top_p < 1.0:
+        logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+
+    # 3. Get probabilities
+    probs = F.softmax(logits, dim=-1)
+
+    # 4. Sample or take argmax
+    if do_sample:
+        curr_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
+    else:
+        curr_tokens = torch.argmax(probs, dim=-1)
+
+    curr_embeds = codebook(curr_tokens)
+
+    return curr_tokens, curr_embeds
+
+
+def flip_tensor_elements_uniform_prob(tensor: torch.Tensor, p_max: float) -> torch.Tensor:
+    if not 0.0 <= p_max <= 1.0:
+        raise ValueError(f"p_max must in [0.0, 1.0]")
+
+    r1 = torch.rand_like(tensor)
+    r2 = torch.rand_like(tensor)
+
+    flip_mask = r1 < p_max * r2
+
+    multiplier = torch.where(flip_mask, -1.0, 1.0)
+    multiplier = multiplier.to(tensor.dtype)
+
+    flipped_tensor = tensor * multiplier
+    return flipped_tensor
+
+def gaussian_sample(raw_output):
+    mu, log_var = raw_output.chunk(2, dim=-1)
+    sigma = torch.exp(0.5 * log_var)
+    sample = mu + torch.randn_like(mu) * sigma
+
+    return sample
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0):
+    """
+    grid_size: int or tuple/list of (h, w)
+    return:
+    pos_embed: [grid_h*grid_w, embed_dim] or [extra_tokens+grid_h*grid_w, embed_dim] (w/ or w/o cls_token)
+    """
+    if isinstance(grid_size, int):
+        grid_h_size, grid_w_size = grid_size, grid_size
+    else:
+        grid_h_size, grid_w_size = grid_size
+
+    grid_h = torch.arange(grid_h_size, dtype=torch.float32) / pe_interpolation
+    grid_w = torch.arange(grid_w_size, dtype=torch.float32) / pe_interpolation
+    
+    grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing='xy')
+    
+    grid = torch.stack([grid_w, grid_h], dim=0) # shape: (2, grid_h_size, grid_w_size)
+
+    grid = grid.reshape([2, 1, grid_h_size, grid_w_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+
+    if cls_token and extra_tokens > 0:
+        pos_embed = torch.cat([torch.zeros([extra_tokens, embed_dim]), pos_embed], dim=0)
+    
+    return pos_embed
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = torch.cat([emb_h, emb_w], dim=1)  # (H*W, D)
+    return emb
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    
+    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    emb = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
+    return emb
+
+def remove_first_user_block(x: str) -> str:
+    start_marker = "<|im_start|>user\n"
+    end_marker = "<|im_end|>\n"
+    start_index = x.find(start_marker)
+    if start_index == -1:
+        return x
+    end_index = x.find(end_marker, start_index + len(start_marker))
+    if end_index == -1:
+        return x
+    result = x[:start_index] + x[end_index + len(end_marker):]
+    return result

modeling/vision_encoder/autoencoder.py

@@ -0,0 +1,520 @@
+import torch
+import torch.nn as nn
+import math
+from einops import rearrange
+import torch.nn.functional as F
+from collections import defaultdict
+
+
+def swish(x):
+    return x*torch.sigmoid(x)
+
+class ResBlock(nn.Module):
+    def __init__(self, 
+                 in_filters,
+                 out_filters,
+                 use_conv_shortcut = False,
+                 use_agn = False,
+                 ) -> None:
+        super().__init__()
+
+        self.in_filters = in_filters
+        self.out_filters = out_filters
+        self.use_conv_shortcut = use_conv_shortcut
+        self.use_agn = use_agn
+
+        if not use_agn: ## agn is GroupNorm likewise skip it if has agn before
+            self.norm1 = nn.GroupNorm(32, in_filters, eps=1e-6)
+        self.norm2 = nn.GroupNorm(32, out_filters, eps=1e-6)
+
+        self.conv1 = nn.Conv2d(in_filters, out_filters, kernel_size=(3, 3), padding=1, bias=False)
+        self.conv2 = nn.Conv2d(out_filters, out_filters, kernel_size=(3, 3), padding=1, bias=False)
+
+        if in_filters != out_filters:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = nn.Conv2d(in_filters, out_filters, kernel_size=(3, 3), padding=1, bias=False)
+            else:
+                self.nin_shortcut = nn.Conv2d(in_filters, out_filters, kernel_size=(1, 1), padding=0, bias=False)
+    
+
+    def forward(self, x, **kwargs):
+        residual = x
+
+        if not self.use_agn:
+            x = self.norm1(x)
+        x = swish(x)
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = swish(x)
+        x = self.conv2(x)
+        if self.in_filters != self.out_filters:
+            if self.use_conv_shortcut:
+                residual = self.conv_shortcut(residual)
+            else:
+                residual = self.nin_shortcut(residual)
+
+        return x + residual
+    
+class Encoder(nn.Module):
+    def __init__(self, *, ch, out_ch, in_channels, num_res_blocks, z_channels, ch_mult=(1, 2, 2, 4), 
+                resolution=None, double_z=False,
+                ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.z_channels = z_channels
+        self.resolution = resolution
+
+        self.num_res_blocks = num_res_blocks
+        self.num_blocks = len(ch_mult)
+        
+        self.conv_in = nn.Conv2d(in_channels,
+                                 ch,
+                                 kernel_size=(3, 3),
+                                 padding=1,
+                                 bias=False
+        )
+
+        ## construct the model
+        self.down = nn.ModuleList()
+
+        in_ch_mult = (1,)+tuple(ch_mult)
+        for i_level in range(self.num_blocks):
+            block = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level] #[1, 1, 2, 2, 4]
+            block_out = ch*ch_mult[i_level] #[1, 2, 2, 4]
+            for _ in range(self.num_res_blocks):
+                block.append(ResBlock(block_in, block_out))
+                block_in = block_out
+            
+            down = nn.Module()
+            down.block = block
+            if i_level < self.num_blocks - 1:
+                down.downsample = nn.Conv2d(block_out, block_out, kernel_size=(3, 3), stride=(2, 2), padding=1)
+
+            self.down.append(down)
+        
+        ### mid
+        self.mid_block = nn.ModuleList()
+        for res_idx in range(self.num_res_blocks):
+            self.mid_block.append(ResBlock(block_in, block_in))
+        
+        ### end
+        self.norm_out = nn.GroupNorm(32, block_out, eps=1e-6)
+        self.conv_out = nn.Conv2d(block_out, z_channels, kernel_size=(1, 1))
+            
+    def forward(self, x):
+
+        ## down
+        x = self.conv_in(x)
+        for i_level in range(self.num_blocks):
+            for i_block in range(self.num_res_blocks):
+                x = self.down[i_level].block[i_block](x)
+            
+            if i_level <  self.num_blocks - 1:
+                x = self.down[i_level].downsample(x)
+        
+        ## mid 
+        for res in range(self.num_res_blocks):
+            x = self.mid_block[res](x)
+        
+
+        x = self.norm_out(x)
+        x = swish(x)
+        x = self.conv_out(x)
+
+        return x
+
+class Decoder(nn.Module):
+    def __init__(self, *, ch, out_ch, in_channels, num_res_blocks, z_channels, ch_mult=(1, 2, 2, 4), 
+                resolution=None, double_z=False,) -> None:
+        super().__init__()
+
+        self.ch = ch
+        self.num_blocks = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        block_in = ch*ch_mult[self.num_blocks-1]
+
+        self.conv_in = nn.Conv2d(
+            z_channels, block_in, kernel_size=(3, 3), padding=1, bias=True
+        )
+
+        self.mid_block = nn.ModuleList()
+        for res_idx in range(self.num_res_blocks):
+            self.mid_block.append(ResBlock(block_in, block_in))
+        
+        self.up = nn.ModuleList()
+
+        self.adaptive = nn.ModuleList()
+
+        for i_level in reversed(range(self.num_blocks)):
+            block = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            self.adaptive.insert(0, AdaptiveGroupNorm(z_channels, block_in))
+            for i_block in range(self.num_res_blocks):
+                block.append(ResBlock(block_in, block_out))
+                block_in = block_out
+            
+            up = nn.Module()
+            up.block = block
+            if i_level > 0:
+                up.upsample = Upsampler(block_in)
+            self.up.insert(0, up)
+        
+        self.norm_out = nn.GroupNorm(32, block_in, eps=1e-6)
+
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=(3, 3), padding=1)
+    
+    def forward(self, z):
+        
+        style = z.clone() #for adaptive groupnorm
+
+        z = self.conv_in(z)
+
+        ## mid
+        for res in range(self.num_res_blocks):
+            z = self.mid_block[res](z)
+        
+        ## upsample
+        for i_level in reversed(range(self.num_blocks)):
+            ### pass in each resblock first adaGN
+            z = self.adaptive[i_level](z, style)
+            for i_block in range(self.num_res_blocks):
+                z = self.up[i_level].block[i_block](z)
+            
+            if i_level > 0:
+                z = self.up[i_level].upsample(z)
+        
+        z = self.norm_out(z)
+        z = swish(z)
+        z = self.conv_out(z)
+
+        return z
+
+def depth_to_space(x: torch.Tensor, block_size: int) -> torch.Tensor:
+    """ Depth-to-Space DCR mode (depth-column-row) core implementation.
+
+        Args:
+            x (torch.Tensor): input tensor. The channels-first (*CHW) layout is supported.
+            block_size (int): block side size
+    """
+    # check inputs
+    if x.dim() < 3:
+        raise ValueError(
+            f"Expecting a channels-first (*CHW) tensor of at least 3 dimensions"
+        )
+    c, h, w = x.shape[-3:]
+
+    s = block_size**2
+    if c % s != 0:
+        raise ValueError(
+            f"Expecting a channels-first (*CHW) tensor with C divisible by {s}, but got C={c} channels"
+        )
+
+    outer_dims = x.shape[:-3]
+
+    # splitting two additional dimensions from the channel dimension
+    x = x.view(-1, block_size, block_size, c // s, h, w)
+
+    # putting the two new dimensions along H and W
+    x = x.permute(0, 3, 4, 1, 5, 2)
+
+    # merging the two new dimensions with H and W
+    x = x.contiguous().view(*outer_dims, c // s, h * block_size,
+                            w * block_size)
+
+    return x
+
+class Upsampler(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out = None
+    ):
+        super().__init__()
+        dim_out = dim * 4
+        self.conv1 = nn.Conv2d(dim, dim_out, (3, 3), padding=1)
+        self.depth2space = depth_to_space
+
+    def forward(self, x):
+        """
+        input_image: [B C H W]
+        """
+        out = self.conv1(x)
+        out = self.depth2space(out, block_size=2)
+        return out
+        
+class AdaptiveGroupNorm(nn.Module):
+    def __init__(self, z_channel, in_filters, num_groups=32, eps=1e-6):
+        super().__init__()
+        self.gn = nn.GroupNorm(num_groups=32, num_channels=in_filters, eps=eps, affine=False)
+        # self.lin = nn.Linear(z_channels, in_filters * 2)
+        self.gamma = nn.Linear(z_channel, in_filters)
+        self.beta = nn.Linear(z_channel, in_filters)
+        self.eps = eps
+    
+    def forward(self, x, quantizer):
+        B, C, _, _ = x.shape
+        # quantizer = F.adaptive_avg_pool2d(quantizer, (1, 1))
+        ### calcuate var for scale
+        scale = rearrange(quantizer, "b c h w -> b c (h w)")
+        scale = scale.var(dim=-1) + self.eps #not unbias
+        scale = scale.sqrt()
+        scale = self.gamma(scale).view(B, C, 1, 1)
+
+        ### calculate mean for bias
+        bias = rearrange(quantizer, "b c h w -> b c (h w)")
+        bias = bias.mean(dim=-1)
+        bias = self.beta(bias).view(B, C, 1, 1)
+       
+        x = self.gn(x)
+        x = scale * x + bias
+
+        return x
+
+class GANDecoder(nn.Module):
+    def __init__(self, *, ch, out_ch, in_channels, num_res_blocks, z_channels, ch_mult=(1, 2, 2, 4), 
+                resolution=None, double_z=False,) -> None:
+        super().__init__()
+
+        self.ch = ch
+        self.num_blocks = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        block_in = ch*ch_mult[self.num_blocks-1]
+
+        self.conv_in = nn.Conv2d(
+            z_channels * 2, block_in, kernel_size=(3, 3), padding=1, bias=True
+        )
+
+        self.mid_block = nn.ModuleList()
+        for res_idx in range(self.num_res_blocks):
+            self.mid_block.append(ResBlock(block_in, block_in))
+        
+        self.up = nn.ModuleList()
+
+        self.adaptive = nn.ModuleList()
+
+        for i_level in reversed(range(self.num_blocks)):
+            block = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            self.adaptive.insert(0, AdaptiveGroupNorm(z_channels, block_in))
+            for i_block in range(self.num_res_blocks):
+                # if i_block == 0:
+                #     block.append(ResBlock(block_in, block_out, use_agn=True))
+                # else:
+                block.append(ResBlock(block_in, block_out))
+                block_in = block_out
+            
+            up = nn.Module()
+            up.block = block
+            if i_level > 0:
+                up.upsample = Upsampler(block_in)
+            self.up.insert(0, up)
+        
+        self.norm_out = nn.GroupNorm(32, block_in, eps=1e-6)
+
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=(3, 3), padding=1)
+    
+    def forward(self, z):
+        
+        style = z.clone() #for adaptive groupnorm
+
+        noise = torch.randn_like(z).to(z.device) #generate noise
+        z = torch.cat([z, noise], dim=1) #concat noise to the style vector
+        z = self.conv_in(z)
+
+        ## mid
+        for res in range(self.num_res_blocks):
+            z = self.mid_block[res](z)
+        
+        ## upsample
+        for i_level in reversed(range(self.num_blocks)):
+            ### pass in each resblock first adaGN
+            z = self.adaptive[i_level](z, style)
+            for i_block in range(self.num_res_blocks):
+                z = self.up[i_level].block[i_block](z)
+            
+            if i_level > 0:
+                z = self.up[i_level].upsample(z)
+        
+        z = self.norm_out(z)
+        z = swish(z)
+        z = self.conv_out(z)
+
+        return z
+    
+
+class VQModel(nn.Module):
+    def __init__(self,
+                ddconfig,
+                checkpoint=None,
+                gan_decoder = False,
+                ):
+        super().__init__()
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = GANDecoder(**ddconfig) if gan_decoder else Decoder(**ddconfig)
+
+        # Load weights from the checkpoint
+        if checkpoint is not None:
+            self.load_from_ckpt(checkpoint)
+
+    def load_from_ckpt(self, checkpoint):
+        state = torch.load(checkpoint, mmap=True, map_location="cpu")
+        log_info = self.load_state_dict(state["state_dict"], strict=False)
+        has_missing_keys = bool(log_info.missing_keys)
+        has_unexpected_keys = bool(log_info.unexpected_keys)
+        if not has_missing_keys:
+            print(f"Successfully loaded all weights from checkpoint: {checkpoint}")
+        else:
+            if has_missing_keys:
+                print("Missing keys (model layers not in checkpoint):")
+                for key in log_info.missing_keys:
+                    print(f"  - {key}")
+        if False and has_unexpected_keys:
+            print("\nUnexpected keys (checkpoint layers not in model):")
+            for key in log_info.unexpected_keys:
+                print(f"  - {key}")
+
+    def encode(self, x):
+        h = self.encoder(x)
+        codebook_value = torch.Tensor([1.0]).to(h)
+        quant_h = torch.where(h > 0, codebook_value, -codebook_value)      # higher than 0 filled 
+
+        return quant_h
+    
+    # def vt_forward(self, image_list):
+    #     q_list = []
+    #     for x in image_list:
+    #         quant = self.encode(x)
+    #         quant = rearrange(quant.squeeze(0), "c h w -> (h w) c")
+    #         q_list.append(quant)
+
+    #     return torch.cat(q_list, dim=0)   
+    
+
+    def vt_forward(self, image_list, max_bs=32, ps=1):
+        groups = defaultdict(list)  # {(H, W): [(idx, image_tensor), ...]}
+        for i, img in enumerate(image_list):
+            _, _, H, W = img.shape
+            groups[(H, W)].append((i, img))
+
+        output = [None] * len(image_list)
+
+        for (H, W), items in groups.items():
+            for start in range(0, len(items), max_bs):
+                chunk = items[start:start + max_bs]
+                idxs = [x[0] for x in chunk]
+                imgs = [x[1] for x in chunk]
+
+                batch = torch.cat(imgs, dim=0)  # [B, 3, H, W]
+
+                quant = self.encode(batch)                      # [B, C, h, w]
+
+                for b in range(quant.size(0)):
+                    q = rearrange(quant[b], "c (h p1) (w p2) -> (h w p1 p2) c", p1=ps, p2=ps)
+                    output[idxs[b]] = q
+
+        return torch.cat(output, dim=0)
+    
+    def vt_forward_maxpad(
+        self,
+        image_list,
+        max_bs=32,
+        stride=32,
+        min_size=256,
+        max_size=2048,
+        max_pixels=1024 * 1024,
+        normal_buckets=(384, 512, 768, 1024),
+    ):
+        """
+        image_list: list of [1, 3, H, W]
+        return: Tensor [(sum_i Hi*Wi/stride^2), C]
+        """
+
+        def is_long_image(H, W):
+            major = max(H, W)
+            minor = min(H, W)
+            return (
+                major >= 1024 and
+                minor <= 768 and
+                major / minor >= 1.5
+            )
+
+        groups = defaultdict(list)
+        sizes = {}
+
+        for idx, img in enumerate(image_list):
+            _, _, H, W = img.shape
+
+            # assert H >= min_size and W >= min_size
+            # assert H <= max_size and W <= max_size
+            # assert H * W <= max_pixels, f"image is too large: {H}x{W}"
+
+            if is_long_image(H, W):
+                bucket = "long"
+            else:
+                major = max(H, W)
+                for b in normal_buckets:
+                    if major <= b:
+                        bucket = b
+                        break
+                else:
+                    bucket = "long"
+
+            groups[bucket].append(idx)
+            sizes[idx] = (H, W)
+
+        output = [None] * len(image_list)
+
+
+        for bucket, idxs in groups.items():
+            imgs = [image_list[i] for i in idxs]
+
+            for start in range(0, len(imgs), max_bs):
+                batch_imgs = imgs[start:start + max_bs]
+                batch_idxs = idxs[start:start + max_bs]
+
+                H_max = max(img.shape[-2] for img in batch_imgs)
+                W_max = max(img.shape[-1] for img in batch_imgs)
+
+                H_pad = math.ceil(H_max / stride) * stride
+                W_pad = math.ceil(W_max / stride) * stride
+
+                padded = []
+                for img in batch_imgs:
+                    _, _, H, W = img.shape
+                    pad_h = H_pad - H
+                    pad_w = W_pad - W
+                    padded.append(F.pad(img, (0, pad_w, 0, pad_h)))
+
+                batch = torch.cat(padded, dim=0)  # [B, 3, H_pad, W_pad]
+
+                quant = self.encode(batch)  # [B, C, h', w']
+
+                for i, q in enumerate(quant):
+                    H, W = sizes[batch_idxs[i]]
+                    h_lat = math.ceil(H / stride)
+                    w_lat = math.ceil(W / stride)
+
+                    q = q[:, :h_lat, :w_lat]
+                    q = rearrange(q, "c h w -> (h w) c")
+
+                    output[batch_idxs[i]] = q
+
+        return torch.cat(output, dim=0)
+
+
+    def decode(self, quant):
+        dec = self.decoder(quant)
+        return dec
+
+    def forward(self, input):
+        quant = self.encode(input)
+        dec = self.decode(quant)
+        return dec, quant

modeling/vision_head/flow_head_parallel_x.py

@@ -0,0 +1,342 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .sampling_x import euler_maruyama
+
+from flash_attn import flash_attn_func
+
+
+def timestep_embedding(t, dim, max_period=10000, time_factor: float = 1000.0):
+    half = dim // 2
+    t = time_factor * t.float()
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+        / half
+    )
+
+    args = t[:, None] * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+def time_shift_func(t: torch.Tensor, flow_shift: float = 1., sigma: float = 1.):
+    return (1 / flow_shift) / ( (1 / flow_shift) + (1 / t - 1) ** sigma)
+
+class DiffHead(nn.Module):
+    def __init__(
+        self,
+        ch_target,
+        ch_cond,
+        ch_latent,
+        depth_latent,
+        depth_adanln,
+        grad_checkpointing=False,
+        time_shift=1.,
+        time_schedule='logit_normal',
+        P_mean=0.,
+        P_std=1.,
+        parallel_num=4,
+        diff_batch_mul=1,
+        use_swiglu=False,
+    ):
+        super(DiffHead, self).__init__()
+        self.ch_target = ch_target
+        self.time_shift = time_shift
+        self.time_schedule = time_schedule
+        self.P_mean = P_mean
+        self.P_std = P_std
+        self.diff_batch_mul = diff_batch_mul
+
+        self.net = TransEncoder(
+            in_channels=ch_target,
+            model_channels=ch_latent,
+            z_channels=ch_cond,
+            num_res_blocks=depth_latent,
+            num_ada_ln_blocks=depth_adanln,
+            grad_checkpointing=grad_checkpointing,
+            parallel_num=parallel_num,
+            use_swiglu=use_swiglu
+        )
+
+    def forward(self, x, cond):
+        with torch.autocast(device_type="cuda", enabled=False):
+            with torch.no_grad():
+                if self.time_schedule == 'logit_normal':
+                    t = (torch.randn((x.shape[0]), device=x.device) * self.P_std + self.P_mean).sigmoid()
+                    if self.time_shift != 1.:
+                        t = time_shift_func(t, self.time_shift)
+                elif self.time_schedule == 'uniform':
+                    t = torch.rand((x.shape[0]), device=x.device)
+                    if self.time_shift != 1.:
+                        t = time_shift_func(t, self.time_shift)
+                else:
+                    raise NotImplementedError(f"unknown time_schedule {self.time_schedule}")
+                e = torch.randn_like(x)
+                ti = t.view(-1, 1, 1)
+                z = (1.0 - ti) * e + ti * x
+                v = (x - z) / (1 - ti).clamp_min(0.05)
+
+        if self.diff_batch_mul > 1:
+            chunks = self.diff_batch_mul
+            x_pred_list = []
+
+            z_chunks = torch.chunk(z, chunks, dim=0)
+            t_chunks = torch.chunk(t, chunks, dim=0) 
+            cond_chunks = torch.chunk(cond, chunks, dim=0)
+            for z_i, t_i, cond_i in zip(z_chunks, t_chunks, cond_chunks):
+                output_i = self.net(z_i, t_i, cond_i)
+                x_pred_list.append(output_i)
+            x_pred = torch.cat(x_pred_list, dim=0)
+        else:
+            x_pred = self.net(z, t, cond)
+            
+        v_pred = (x_pred - z) / (1 - ti).clamp_min(0.05)
+
+        with torch.autocast(device_type="cuda", enabled=False):
+            v_pred = v_pred.float()
+            loss = torch.mean((v - v_pred) ** 2, dim=2)
+        return loss
+
+    def sample(
+        self,
+        z,
+        cfg,
+        num_sampling_steps,
+    ):
+        return euler_maruyama(
+            self.ch_target,
+            self.net.forward,
+            z,
+            cfg,
+            num_sampling_steps=num_sampling_steps,
+            time_shift = self.time_shift,
+        )
+
+    def initialize_weights(self):
+        self.net.initialize_weights()
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    def forward(self, t):
+        t_freq = timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class ResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.norm = nn.LayerNorm(channels, eps=1e-6, elementwise_affine=True)
+        hidden_dim = int(channels * 1.5)
+        self.w1 = nn.Linear(channels, hidden_dim * 2, bias=True)
+        self.w2 = nn.Linear(hidden_dim, channels, bias=True)
+
+    def forward(self, x, scale, shift, gate):
+        h = self.norm(x) * (1 + scale) + shift
+        h1, h2 = self.w1(h).chunk(2, dim=-1)
+        h = self.w2(F.silu(h1) * h2)
+        return x + h * gate
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, channels, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(channels, eps=1e-6, elementwise_affine=False)
+        self.ada_ln_modulation = nn.Linear(channels, channels * 2, bias=True)
+        self.linear = nn.Linear(channels, out_channels, bias=True)
+
+    def forward(self, x, y):
+        scale, shift = self.ada_ln_modulation(y).chunk(2, dim=-1)
+        x = self.norm_final(x) * (1.0 + scale) + shift
+        x = self.linear(x)
+        return x
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        n_head,
+    ):
+        super().__init__()
+        assert dim % n_head == 0
+        self.dim = dim
+        self.head_dim = dim // n_head
+        self.scale = self.head_dim**-0.5
+        self.n_head = n_head
+        total_kv_dim = (self.n_head * 3) * self.head_dim
+
+        self.wqkv = nn.Linear(dim, total_kv_dim, bias=True)
+        self.wo = nn.Linear(dim, dim, bias=True)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        bsz, seqlen, _ = x.shape
+        xq, xk, xv = self.wqkv(x).chunk(3, dim=-1)
+
+        xq = xq.view(bsz, seqlen, self.n_head, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_head, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_head, self.head_dim)
+
+        if seqlen <= 32:
+            xq, xk, xv = map(lambda x: x.transpose(1, 2), (xq, xk, xv))
+            xq = xq * self.scale
+            attn = xq @ xk.transpose(-1, -2)
+            attn = F.softmax(attn, dim=-1)
+            output = (attn @ xv).transpose(1, 2).contiguous()
+        else:
+            output = flash_attn_func(
+                    xq,
+                    xk,
+                    xv,
+                    causal=False,
+                )
+
+        output = output.view(bsz, seqlen, self.dim)
+
+        output = self.wo(output)
+        return output
+
+class TransBlock(nn.Module):
+    def __init__(self, channels, use_swiglu=False):
+        super().__init__()
+        self.channels = channels
+        self.norm1 = nn.LayerNorm(channels, eps=1e-6, elementwise_affine=True)
+        self.attn = Attention(channels, n_head=channels // 128)
+
+        self.norm2 = nn.LayerNorm(channels, eps=1e-6, elementwise_affine=True)
+        hidden_dim = int(channels * 1.5)
+        self.use_swiglu = use_swiglu
+        if not self.use_swiglu:
+            self.mlp = nn.Sequential(
+                nn.Linear(self.channels, hidden_dim),
+                nn.SiLU(),
+                nn.Linear(hidden_dim, self.channels),
+            )
+        else:
+            self.w1 = nn.Linear(channels, hidden_dim * 2, bias=True)
+            self.w2 = nn.Linear(hidden_dim, channels, bias=True)
+
+    def forward(self, x, scale1, shift1, gate1, scale2, shift2, gate2):
+        h = self.norm1(x) * (1 + scale1) + shift1
+        h = self.attn(h)
+        x = x + h * gate1
+        h = self.norm2(x) * (1 + scale2) + shift2
+        if not self.use_swiglu:
+            h = self.mlp(h)
+        else:
+            h1, h2 = self.w1(h).chunk(2, dim=-1)
+            h = self.w2(F.silu(h1) * h2)
+        return x + h * gate2
+
+class TransEncoder(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        z_channels,
+        num_res_blocks,
+        num_ada_ln_blocks=2,
+        grad_checkpointing=False,
+        parallel_num=4,
+        use_swiglu=False,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = in_channels
+        self.num_res_blocks = num_res_blocks
+        self.grad_checkpointing = grad_checkpointing
+        self.parallel_num = parallel_num
+
+        self.time_embed = TimestepEmbedder(model_channels)
+        self.cond_embed = nn.Linear(z_channels, model_channels)
+
+        self.input_proj = nn.Linear(in_channels, model_channels)
+        self.res_blocks = nn.ModuleList()
+        for i in range(num_res_blocks):
+            self.res_blocks.append(
+                TransBlock(
+                    model_channels,
+                    use_swiglu
+                )
+            )
+        # share adaLN for consecutive blocks, to save computation and parameters
+        self.ada_ln_blocks = nn.ModuleList()
+        for i in range(num_ada_ln_blocks):
+            self.ada_ln_blocks.append(
+                nn.Linear(model_channels, model_channels * 6, bias=True)
+            )
+        self.ada_ln_switch_freq = max(1, num_res_blocks // num_ada_ln_blocks)
+        assert (
+            num_res_blocks % self.ada_ln_switch_freq
+        ) == 0, "num_res_blocks must be divisible by num_ada_ln_blocks"
+        self.final_layer = FinalLayer(model_channels, self.out_channels)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP
+        nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
+
+        for block in self.ada_ln_blocks:
+            nn.init.constant_(block.weight, 0)
+            nn.init.constant_(block.bias, 0)
+
+        # Zero-out output layers
+        nn.init.constant_(self.final_layer.ada_ln_modulation.weight, 0)
+        nn.init.constant_(self.final_layer.ada_ln_modulation.bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def forward(self, x, t, c):
+        x = self.input_proj(x)
+        t = self.time_embed(t).unsqueeze(1)
+        c = self.cond_embed(c)
+
+        y = F.silu(t + c)
+        scale1, shift1, gate1, scale2, shift2, gate2 = self.ada_ln_blocks[0](y).chunk(6, dim=-1)
+
+        for i, block in enumerate(self.res_blocks):
+            if i > 0 and i % self.ada_ln_switch_freq == 0:
+                ada_ln_block = self.ada_ln_blocks[i // self.ada_ln_switch_freq]
+                scale1, shift1, gate1, scale2, shift2, gate2 = ada_ln_block(y).chunk(6, dim=-1)
+            x = block(x, scale1, shift1, gate1, scale2, shift2, gate2)
+
+        output = self.final_layer(x, y)
+
+        # use sigmoid to map to [-1, 1]
+        return 2 * torch.sigmoid(output) - 1

modeling/vision_head/sampling_x.py

@@ -0,0 +1,125 @@
+import torch
+
+def time_shift_func(t: torch.Tensor, flow_shift: float = 1., sigma: float = 1.):
+    return (1 / flow_shift) / ( (1 / flow_shift) + (1 / t - 1) ** sigma)
+
+def get_score_from_velocity(velocity, x, t):
+    alpha_t, d_alpha_t = t, 1
+    sigma_t, d_sigma_t = 1 - t, -1
+    mean = x
+    reverse_alpha_ratio = alpha_t / d_alpha_t
+    var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+    score = (reverse_alpha_ratio * velocity - mean) / var
+    return score
+
+
+def get_velocity_from_cfg(velocity, cfg, cfg_mult):
+    if cfg_mult == 2:
+        cond_v, uncond_v = torch.chunk(velocity, 2, dim=0)
+        velocity = uncond_v + cfg * (cond_v - uncond_v)
+    return velocity
+
+
+# @torch.compile()
+def euler_step(x, v, dt: float, cfg: float, cfg_mult: int):
+    with torch.amp.autocast("cuda", enabled=False):
+        v = v.to(torch.float32)
+        v = get_velocity_from_cfg(v, cfg, cfg_mult)
+        x = x + v * dt
+    return x
+
+
+# @torch.compile()
+def euler_maruyama_step(x, v, t, dt: float, cfg: float, cfg_mult: int):
+    with torch.amp.autocast("cuda", enabled=False):
+        v = v.to(torch.float32)
+        v = get_velocity_from_cfg(v, cfg, cfg_mult)
+        score = get_score_from_velocity(v, x, t)
+        drift = v + (1 - t) * score
+        noise_scale = (2.0 * (1.0 - t) * dt) ** 0.5
+        x = x + drift * dt + noise_scale * torch.randn_like(x)
+    return x
+
+
+def euler_maruyama(
+    input_dim,
+    forward_fn,
+    c: torch.Tensor,
+    cfg: float = 1.0,
+    num_sampling_steps: int = 20,
+    last_step_size: float = 0.05,
+    time_shift: float = 1.,
+):
+    cfg_mult = 1
+    if cfg > 1.0:
+        cfg_mult += 1
+
+    x_shape = list(c.shape)
+    x_shape[0] = x_shape[0] // cfg_mult
+    x_shape[-1] = input_dim
+    x = torch.randn(x_shape, device=c.device)
+    # an = (1.0 - last_step_size) / num_sampling_steps
+    t_all = torch.linspace(0, 1-last_step_size, num_sampling_steps+1, device=c.device, dtype=torch.float32)
+    t_all = time_shift_func(t_all, time_shift)
+    dt = t_all[1:] - t_all[:-1]
+    t = torch.tensor(
+        0.0, device=c.device, dtype=torch.float32
+    )  # use tensor to avoid compile warning
+    t_batch = torch.zeros(c.shape[0], device=c.device)
+    for i in range(num_sampling_steps):
+        t_batch[:] = t
+        combined = torch.cat([x] * cfg_mult, dim=0)
+        output = forward_fn(
+            combined,
+            t_batch,
+            c,
+        )
+        if output.dim() == 2:
+            v = (output - combined) / (1 - t_batch.view(-1,1)).clamp_min(0.05)
+        elif output.dim() == 3:
+            v = (output - combined) / (1 - t_batch.view(-1,1,1)).clamp_min(0.05)
+        x = euler_maruyama_step(x, v, t, dt[i], cfg, cfg_mult)
+        t += dt[i]
+
+    combined = torch.cat([x] * cfg_mult, dim=0)
+    t_batch[:] = 1 - last_step_size
+    output = forward_fn(
+        combined,
+        t_batch,
+        c,
+    )
+    if output.dim() == 2:
+        v = (output - combined) / (1 - t_batch.view(-1,1)).clamp_min(0.05)
+    elif output.dim() == 3:
+        v = (output - combined) / (1 - t_batch.view(-1,1,1)).clamp_min(0.05)
+    x = euler_step(x, v, last_step_size, cfg, cfg_mult)
+
+    return torch.cat([x] * cfg_mult, dim=0)
+
+
+def euler(
+    input_dim,
+    forward_fn,
+    c,
+    cfg: float = 1.0,
+    num_sampling_steps: int = 50,
+):
+    cfg_mult = 1
+    if cfg > 1.0:
+        cfg_mult = 2
+
+    x_shape = list(c.shape)
+    x_shape[0] = x_shape[0] // cfg_mult
+    x_shape[-1] = input_dim
+    x = torch.randn(x_shape, device=c.device)
+    dt = 1.0 / num_sampling_steps
+    t = 0
+    t_batch = torch.zeros(c.shape[0], device=c.device)
+    for _ in range(num_sampling_steps):
+        t_batch[:] = t
+        combined = torch.cat([x] * cfg_mult, dim=0)
+        v = forward_fn(combined, t_batch, c)
+        x = euler_step(x, v, dt, cfg, cfg_mult)
+        t += dt
+
+    return torch.cat([x] * cfg_mult, dim=0)

requirements.txt

@@ -0,0 +1,10 @@
+transformers==4.57.0
+omegaconf
+liger-kernel
+numpy==1.26.4
+huggingface_hub==0.34.4
+einops==0.6.1
+torch==2.7.1
+torchvision==0.22.1
+safetensors==0.6.2
+gradio