Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,15 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples/10000.jpg filter=lfs diff=lfs merge=lfs -text
+samples/10500.jpg filter=lfs diff=lfs merge=lfs -text
+samples/11000.jpg filter=lfs diff=lfs merge=lfs -text
+samples/11500.jpg filter=lfs diff=lfs merge=lfs -text
+samples/12000.jpg filter=lfs diff=lfs merge=lfs -text
+samples/12500.jpg filter=lfs diff=lfs merge=lfs -text
+samples/13000.jpg filter=lfs diff=lfs merge=lfs -text
+samples/7500.jpg filter=lfs diff=lfs merge=lfs -text
+samples/8000.jpg filter=lfs diff=lfs merge=lfs -text
+samples/8500.jpg filter=lfs diff=lfs merge=lfs -text
+samples/9000.jpg filter=lfs diff=lfs merge=lfs -text
+samples/9500.jpg filter=lfs diff=lfs merge=lfs -text

ckpt_0.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dded894389248d371d9ad72ea2bf90e7a5677fc5ec3801e594eeeec200512993
+size 202217053

simplefsq32.py

@@ -0,0 +1,682 @@
+"""
+最简 FSQ 训练 Demo（支持分桶 Bucketing）
+单文件，自包含，可直接运行。
+
+依赖：
+    pip install torch torchvision einops
+
+运行：
+    # 用不同分辨率的 dummy 图片测试分桶
+    python fsq_demo.py --dummy --epochs 2
+
+    # 用真实 WebP 文件夹
+    python fsq_demo.py --data /path/to/webp_folder --epochs 10
+"""
+
+import argparse
+import os
+import math
+import glob
+import random
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset, Sampler
+from torchvision import transforms
+from torchvision.utils import save_image, make_grid
+from PIL import Image
+from einops import rearrange, pack, unpack
+
+import json
+from concurrent.futures import ThreadPoolExecutor
+# ============================================================
+# 1. FSQ 量化器
+# ============================================================
+
+def round_ste(z: torch.Tensor) -> torch.Tensor:
+    """取整 + 直通梯度估计（Straight-Through Estimator）"""
+    zhat = z.round()
+    return z + (zhat - z).detach()
+
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+
+class FSQ(nn.Module):
+    """Finite Scalar Quantization，来自 https://arxiv.org/abs/2309.15505"""
+
+    def __init__(self, levels: List[int]):
+        super().__init__()
+        _levels = torch.tensor(levels, dtype=torch.int32)
+        self.register_buffer("_levels", _levels, persistent=False)
+
+        _basis = torch.cumprod(torch.tensor([1] + levels[:-1]), dim=0, dtype=torch.int32)
+        self.register_buffer("_basis", _basis, persistent=False)
+
+        self.codebook_dim = len(levels)
+        self.codebook_size = int(_levels.prod().item())
+
+    def bound(self, z: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:
+        half_l = (self._levels - 1) * (1 - eps) / 2
+        offset = torch.where(self._levels % 2 == 0, 0.5, 0.0)
+        shift = (offset / half_l).tan()
+        return (z + shift).tanh() * half_l - offset
+
+    def quantize(self, z: torch.Tensor) -> torch.Tensor:
+        quantized = round_ste(self.bound(z))
+        half_width = self._levels // 2
+        return quantized / half_width
+
+    def _scale_and_shift(self, zhat_normalized):
+        half_width = self._levels // 2
+        return (zhat_normalized * half_width) + half_width
+
+    def _scale_and_shift_inverse(self, zhat):
+        half_width = self._levels // 2
+        return (zhat - half_width) / half_width
+
+    def codes_to_indices(self, zhat: torch.Tensor) -> torch.Tensor:
+        assert zhat.shape[-1] == self.codebook_dim
+        zhat = self._scale_and_shift(zhat)
+        return (zhat * self._basis).sum(dim=-1).to(torch.int32)
+
+    def indices_to_codes(self, indices: torch.Tensor) -> torch.Tensor:
+        indices = rearrange(indices, '... -> ... 1')
+        codes_non_centered = (indices // self._basis) % self._levels
+        codes = self._scale_and_shift_inverse(codes_non_centered)
+        return codes
+
+    def forward(self, z: torch.Tensor):
+        is_img = z.ndim >= 4
+        if is_img:
+            z = rearrange(z, 'b d ... -> b ... d')
+            z, ps = pack_one(z, 'b * d')
+
+        assert z.shape[-1] == self.codebook_dim
+
+        codes = self.quantize(z)
+        indices = self.codes_to_indices(codes)
+
+        if is_img:
+            codes = unpack_one(codes, ps, 'b * d')
+            codes = rearrange(codes, 'b ... d -> b d ...')
+            indices = unpack_one(indices, ps, 'b *')
+
+        return codes, indices
+
+
+# ============================================================
+# 2. Encoder / Decoder（16× 下采样，即 4 层 stride=2）
+# ============================================================
+
+class Encoder(nn.Module):
+    def __init__(self, in_channel=3, channel=128, embed_dim=4):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_channel, channel, 4, stride=2, padding=1),   # 256→128
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel, channel, 4, stride=2, padding=1),       # 128→64
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel, channel, 4, stride=2, padding=1),       # 64→32
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel, channel, 4, stride=2, padding=1),       # 32→16
+            nn.ReLU(inplace=True),
+            #nn.Conv2d(channel, channel, 4, stride=2, padding=1),       # 16→8
+            #nn.ReLU(inplace=True),
+            nn.Conv2d(channel, embed_dim, 1),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Decoder(nn.Module):
+    def __init__(self, embed_dim=4, channel=128, out_channel=3):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.ConvTranspose2d(embed_dim, channel, 4, stride=2, padding=1),   # 8→16
+            #nn.ReLU(inplace=True),
+            #nn.ConvTranspose2d(channel, channel, 4, stride=2, padding=1),   # 16→32
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(channel, channel, 4, stride=2, padding=1),     # 32→64
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(channel, channel, 4, stride=2, padding=1),     # 64→128
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(channel, channel, 4, stride=2, padding=1),     # 128→256
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel, out_channel, 1),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# ============================================================
+# 3. VQVAE (FSQ 版)
+# ============================================================
+
+class FSQVAE(nn.Module):
+    def __init__(self, levels: List[int], channel=128, in_channel=3):
+        super().__init__()
+        embed_dim = len(levels)
+        self.encoder = Encoder(in_channel, channel, embed_dim)
+        self.quantizer = FSQ(levels)
+        self.decoder = Decoder(embed_dim, channel, in_channel)
+
+    def forward(self, x, return_id=False):
+        z = self.encoder(x)
+        quant, indices = self.quantizer(z)
+        recon = self.decoder(quant)
+        if return_id:
+            return recon, indices
+        return recon
+
+    @torch.no_grad()
+    def encode(self, x):
+        z = self.encoder(x)
+        _, indices = self.quantizer(z)
+        return indices
+
+    @torch.no_grad()
+    def decode_code(self, indices):
+        codes = self.quantizer.indices_to_codes(indices)
+        codes = rearrange(codes, 'b h w d -> b d h w')
+        return self.decoder(codes)
+
+
+# ============================================================
+# 4. 分桶（Bucketing）
+# ============================================================
+
+def make_buckets(base=256, align=16):
+    """
+    生成一组总像素数接近 base² 的桶，宽高比覆盖常见情况。
+    每个桶的 H、W 都是 align 的倍数（方便下采样 16×）。
+    """
+    aspect_ratios = [0.65, 0.68, 0.72, 0.75, 0.80, 1.0, 1.33, 1.45, 1.78]
+    buckets = set()
+    for ar in aspect_ratios:
+        h = int(round(math.sqrt(base * base * ar)))
+        w = int(round(math.sqrt(base * base / ar)))
+        h = max(align, h // align * align)
+        w = max(align, w // align * align)
+        buckets.add((h, w))
+    return sorted(buckets)
+
+
+def find_bucket(w, h, buckets):
+    """
+    选裁剪面积最小的桶。
+    """
+    best, best_loss = None, float('inf')
+    for (bh, bw) in buckets:
+        scale = max(bw / w, bh / h)
+        new_w = math.ceil(w * scale)
+        new_h = math.ceil(h * scale)
+        # 裁剪掉的面积比例
+        crop_loss = 1 - (bw * bh) / (new_w * new_h)
+        if crop_loss < best_loss:
+            best_loss, best = crop_loss, (bh, bw)
+    return best
+
+
+class BucketImageFolder(Dataset):
+    def __init__(self, root, buckets, recursive=True,
+                 num_threads=2, cache_file="./sizes_cache.json"):
+        self.buckets = buckets
+        exts = ('*.webp', '*.jpg', '*.jpeg', '*.png', '*.bmp')
+        self.files = []
+        for ext in exts:
+            pattern = f"**/{ext}" if recursive else ext
+            self.files += glob.glob(os.path.join(root, pattern), recursive=recursive)
+        self.files = sorted(self.files)
+        if len(self.files) == 0:
+            raise RuntimeError(f"在 {root} 下没找到任何图片")
+
+        print(f"找到 {len(self.files)} 张图片")
+
+        # 扫描尺寸（多线程 + JSON 缓存）
+        sizes = self._scan_sizes(num_threads, cache_file)
+
+        # 分桶
+        self.bucket_map = {b: [] for b in buckets}
+        self.idx_to_bucket = {}
+        for idx, f in enumerate(self.files):
+            if f not in sizes:
+                continue
+            w, h = sizes[f]
+            b = find_bucket(w, h, buckets)
+            self.bucket_map[b].append(idx)
+            self.idx_to_bucket[idx] = b
+
+        # 打印桶分布
+        for b, idxs in self.bucket_map.items():
+            print(f"  桶 {b}: {len(idxs)} 张")
+
+        self.normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+
+    def _scan_sizes(self, num_threads, cache_file):
+        # 1. 尝试从缓存加载
+        if cache_file and os.path.exists(cache_file):
+            print(f"从缓存加载尺寸: {cache_file}")
+            with open(cache_file) as f:
+                cache = json.load(f)
+            if set(cache.keys()) == set(self.files):
+                print(f"缓存有效，{len(cache)} 个尺寸")
+                return {f: tuple(s) for f, s in cache.items()}
+            print("缓存失效，重新扫描")
+
+        # 2. 多线程扫描
+        print(f"多线程扫描 {len(self.files)} 张图（{num_threads} 线程）...")
+
+        def get_size(f):
+            try:
+                with Image.open(f) as im:
+                    return f, im.size
+            except Exception:
+                return f, None
+
+        with ThreadPoolExecutor(num_threads) as ex:
+            results = list(ex.map(get_size, self.files))
+
+        sizes = {f: s for f, s in results if s is not None}
+        print(f"扫描完成: {len(sizes)} 张成功")
+
+        # 3. 保存缓存
+        if cache_file:
+            print(f"保存缓存: {cache_file}")
+            with open(cache_file, "w") as f:
+                json.dump({f: list(s) for f, s in sizes.items()}, f)
+            print(f"缓存已保存")
+
+        return sizes
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, idx):
+        f = self.files[idx]
+        bh, bw = self.idx_to_bucket.get(idx, (256, 256))
+        try:
+            img = Image.open(f)
+            if getattr(img, "is_animated", False):
+                img.seek(0)
+            img = img.convert("RGB")
+            w, h = img.size
+            scale = max(bw / w, bh / h)
+            new_w = max(bw, math.ceil(w * scale))
+            new_h = max(bh, math.ceil(h * scale))
+            img = img.resize((new_w, new_h), Image.BICUBIC)
+            left = (new_w - bw) // 2
+            top = (new_h - bh) // 2
+            img = img.crop((left, top, left + bw, top + bh))
+        except Exception as e:
+            print(f"读取失败 {f}: {e}")
+            img = Image.new("RGB", (bw, bh), (0, 0, 0))
+
+        img = transforms.functional.to_tensor(img)
+        img = self.normalize(img)
+        return img, 0
+
+
+class BucketBatchSampler(Sampler):
+    """
+    分桶 batch 采样器：每个 batch 内所有图来自同一个桶（尺寸一致）。
+    每个 epoch 重新打乱。
+    """
+    def __init__(self, dataset, batch_size, shuffle=True, drop_last=True):
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+
+    def _make_batches(self):
+        batches = []
+        for bucket, indices in self.dataset.bucket_map.items():
+            indices = list(indices)
+            if self.shuffle:
+                random.shuffle(indices)
+            for i in range(0, len(indices), self.batch_size):
+                batch = indices[i:i + self.batch_size]
+                if len(batch) == self.batch_size or not self.drop_last:
+                    batches.append(batch)
+        if self.shuffle:
+            random.shuffle(batches)
+        return batches
+
+    def __iter__(self):
+        return iter(self._make_batches())
+
+    def __len__(self):
+        # 估算
+        total = 0
+        for indices in self.dataset.bucket_map.values():
+            n = len(indices)
+            total += n // self.batch_size if self.drop_last else math.ceil(n / self.batch_size)
+        return total
+
+
+# ============================================================
+# 5. Dummy 数据集（不同分辨率，测试分桶）
+# ============================================================
+
+class DummyBucketDataset(Dataset):
+    """
+    生成不同分辨率的随机图片，用于测试分桶逻辑。
+    每张图随机选一个宽高比。
+    """
+    def __init__(self, size=200, buckets=None):
+        self.size = size
+        self.buckets = buckets
+        # 随机分配每张图到某个桶
+        self.assignments = []
+        for i in range(size):
+            b = random.choice(buckets)
+            self.assignments.append(b)
+
+        # 构建 bucket_map
+        self.bucket_map = {b: [] for b in buckets}
+        for idx, b in enumerate(self.assignments):
+            self.bucket_map[b].append(idx)
+
+        print(f"Dummy 数据集: {size} 张图")
+        for b, idxs in self.bucket_map.items():
+            print(f"  桶 {b}: {len(idxs)} 张")
+
+        self.normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, idx):
+        bh, bw = self.assignments[idx]
+        # 随机图，归一化到 [-1,1]
+        img = torch.randn(3, bh, bw) * 0.5
+        return img, 0
+
+
+# ============================================================
+# 6. 数据加载器
+# ============================================================
+def get_dataloader(args):
+    buckets = make_buckets(base=args.bucket_base, align=16)
+    print(f"桶列表: {buckets}")
+
+    if args.dummy:
+        ds = DummyBucketDataset(size=args.dummy_size, buckets=buckets)
+    else:
+        ds = BucketImageFolder(
+            args.data, buckets,
+            num_threads=2,
+            cache_file="./sizes_cache.json",   # ← 关键
+        )
+
+    sampler = BucketBatchSampler(
+        ds, batch_size=args.batch_size, shuffle=True, drop_last=True
+    )
+
+    return DataLoader(
+        ds,
+        batch_sampler=sampler,
+        num_workers=args.num_workers,
+        pin_memory=torch.cuda.is_available(),
+    )
+
+# ============================================================
+# 7. 感知损失（可选）
+# ============================================================
+
+class SimplePerceptualLoss(nn.Module):
+    """
+    轻量版感知损失：用 torchvision 的 VGG16 前几层特征。
+    """
+    def __init__(self):
+        super().__init__()
+        from torchvision import models
+        vgg = models.vgg16(weights=models.VGG16_Weights.DEFAULT).features
+        self.slice = nn.Sequential(*[vgg[i] for i in range(9)])
+        for p in self.parameters():
+            p.requires_grad = False
+        self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+    def forward(self, x, y):
+        x = (x + 1) / 2
+        y = (y + 1) / 2
+        x = (x - self.mean) / self.std
+        y = (y - self.mean) / self.std
+        fx = self.slice(x)
+        fy = self.slice(y)
+        return F.l1_loss(fx, fy)
+    
+import lpips
+
+class PerceptualLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.lpips = lpips.LPIPS(net='squeeze').cuda().eval()
+        for p in self.lpips.parameters():
+            p.requires_grad = False
+
+    def forward(self, x, y):
+        # 输入 [-1,1]，直接传
+        return self.lpips(x, y).mean()
+    
+def load_checkpointsimpe(model, optimizer, path, device='cuda'):
+    ckpt = torch.load(path, map_location=device)
+
+    model.load_state_dict(ckpt['model_state_dict'])
+    optimizer.load_state_dict(ckpt['optimizer_state_dict'])
+
+    start_epoch = ckpt['epoch'] + 1
+    print(f"从 epoch {ckpt['epoch']} 恢复，从 epoch {start_epoch} 继续")
+
+    return start_epoch
+
+def load_checkpoint(model, optimizer, scheduler, scaler, path, device='cuda'):
+    ckpt = torch.load(path, map_location=device)
+
+    model.load_state_dict(ckpt['model_state_dict'])
+    optimizer.load_state_dict(ckpt['optimizer_state_dict'])
+
+    if 'scheduler_state_dict' in ckpt and scheduler is not None:
+        scheduler.load_state_dict(ckpt['scheduler_state_dict'])
+
+    if 'scaler_state_dict' in ckpt and scaler is not None:
+        scaler.load_state_dict(ckpt['scaler_state_dict'])
+
+    start_epoch = ckpt['epoch'] + 1
+    num_iter = ckpt.get('num_iter', 0)
+    print(f"恢复: epoch {ckpt['epoch']} → {start_epoch}, iter {num_iter}")
+
+    return start_epoch, num_iter
+
+# ============================================================
+# 8. 训练
+# ============================================================
+
+import time
+import bitsandbytes as bnb
+def train(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    os.makedirs(args.out_dir, exist_ok=True)
+    os.makedirs(os.path.join(args.out_dir, 'samples'), exist_ok=True)
+
+    loader = get_dataloader(args)
+
+    model = FSQVAE(levels=args.levels, channel=args.channel).to(device)
+    #if args.resume:
+        #start_epoch = load_checkpoint(model, optimizer, args.resume, device)
+    #加载预训练
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"模型参数量: {n_params/1e6:.2f} M")
+    print(f"FSQ levels: {args.levels}, 码本大小: {int(torch.tensor(args.levels).prod())}")
+    print(f"下采样倍率: 16×（4 层 stride=2）")
+
+    #optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    optimizer = bnb.optim.Adam8bit(model.parameters(), lr=args.lr)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=args.epochs * len(loader)
+    )
+
+    percep = None
+    if args.perceptual_weight > 0:
+        #percep = SimplePerceptualLoss().to(device).eval()
+        percep = PerceptualLoss().to(device).eval()
+        print("启用感知损失")
+
+    l1_fn = nn.L1Loss()
+    num_iter = 0
+
+    # 混合精度
+    scaler = torch.cuda.amp.GradScaler()
+
+    if args.resume:
+        start_epoch, num_iter = load_checkpoint(model, optimizer, scheduler, scaler, args.resume, device)
+        #start_epoch = load_checkpointsimpe(model, optimizer, args.resume, device)
+    for epoch in range(args.epochs):
+        model.train()
+        for imgs, _ in loader:
+            imgs = imgs.to(device)
+            num_iter += 1
+
+            #t1 = time.time()
+            with torch.cuda.amp.autocast():
+                recon = model(imgs)
+                l1 = l1_fn(imgs, recon)
+            loss = l1
+            #t2 = time.time()
+            #print(f"[前向] {t2-t1:.3f}s | 输入 {tuple(imgs.shape)}")
+
+            if percep is not None:
+                p_loss = percep(imgs, recon)
+                loss = loss + args.perceptual_weight * p_loss
+            #t3 = time.time()
+            #print(f"[感知] {t3-t2:.3f}s")
+
+            optimizer.zero_grad()
+            #loss.backward()
+            scaler.scale(loss).backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            #optimizer.step()
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+            #t4 = time.time()
+            #print(f"[反向] {t4-t3:.3f}s")
+
+            if num_iter % args.log_interval == 0:
+                msg = f"epoch {epoch} | iter {num_iter} | l1 {l1.item():.4f}"
+                if percep is not None:
+                    msg += f" | percep {p_loss.item():.4f}"
+                print(msg)
+
+            if num_iter % args.sample_interval == 0:
+                with torch.no_grad(): #图片采样
+                    n = min(8, imgs.shape[0])
+                    grid = make_grid( 
+                        torch.cat([imgs[:n], recon[:n]]),
+                        nrow=n,
+                    )
+                save_image(
+                    grid,
+                    os.path.join(args.out_dir, 'samples', f'{num_iter}.jpg'),
+                    normalize=True,
+                )
+
+            if num_iter % args.save_interval == 0 :
+               # ---- 保存（含完整状态）----
+                torch.save({
+                    'epoch': epoch,
+                    'num_iter': num_iter,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'scheduler_state_dict': scheduler.state_dict(),
+                    'scaler_state_dict': scaler.state_dict(),
+                }, os.path.join(args.out_dir, f'ckpt_{epoch}.pt'))
+
+
+# ============================================================
+# 9. 参数
+# ============================================================
+###改用bnb8bit训练
+def get_args():
+    p = argparse.ArgumentParser()
+    # 数据
+    p.add_argument('--data', type=str, default=None, help='图片文件夹路径')
+    p.add_argument('--dummy', action='store_true', help='用随机数据测试')
+    p.add_argument('--dummy-size', type=int, default=200)
+    p.add_argument('--num-workers', type=int, default=1)
+    p.add_argument('--bucket-base', type=int, default=256,
+                   help='桶基准尺寸，桶的总像素数接近 base²')
+    # 模型
+    p.add_argument('--levels', type=int, nargs='+', default=[8, 5, 5, 5],
+                   help='FSQ levels，例如 8 5 5 5')
+    p.add_argument('--channel', type=int, default=128)
+    # 训练
+    p.add_argument('--epochs', type=int, default=10)
+    p.add_argument('--batch-size', type=int, default=32)
+    p.add_argument('--lr', type=float, default=1e-4)
+    p.add_argument('--perceptual-weight', type=float, default=0.1)
+    # 日志
+    p.add_argument('--log-interval', type=int, default=5)
+    p.add_argument('--sample-interval', type=int, default=500)
+    p.add_argument('--save-interval', type=int, default=1)
+    p.add_argument('--out-dir', type=str, default='./fsq_demo_out')
+
+    p.add_argument('--resume', type=str, default=None,
+                    help='检查点路径，如 ./fsqsmall_out/ckpt_0.pt；None 表示从头训练')
+    return p.parse_args()
+
+
+
+
+####
+
+'''
+#小模型配置 1mb
+!python simplefsq32.py \
+    --data output_256/ \
+    --epochs 10 \
+    --batch-size 64 \
+    --channel 128 \
+    --levels 8 5 5 5 \
+    --bucket-base 256 \
+    --lr 1e-4 \
+    --perceptual-weight 0.1 \
+    --log-interval 50 \
+    --sample-interval 500 \
+    --save-interval 500 \
+    --out-dir ./fsqsmall_out
+
+#中等 6mb模型测试
+
+!python simplefsq1.py \
+    --data output_256/ \
+    --epochs 1 \
+    --batch-size 8 \
+    --channel 256 \
+    --levels 7 5 5 5 5 \
+    --bucket-base 256 \
+    --lr 1e-4 \
+    --perceptual-weight 0.1 \
+    --log-interval 50 \
+    --sample-interval 500 \
+    --save-interval 500 \
+    --out-dir ./fsq_out
+'''
+# ============================================================
+# 10. 入口
+# ============================================================
+
+if __name__ == '__main__':
+    args = get_args()
+    print(args)
+    train(args)