+
+
+[](https://huggingface.co/papers/2602.02493)
+[](https://arxiv.org/abs/2602.02493)
+[](https://zehong-ma.github.io/PixelGen/)
+[](https://dd0d187fc54e4b00ee.gradio.live/)
+[](https://github.com/Zehong-Ma/PixelGen/)
+
+
+
+ 
+ Figure 1: Visualization of 512x512 images generated by our PixelGen. +
+
+
+## 🫖 Introduction
+We introduce PixelGen, a simple pixel diffusion framework with perceptual loss. Instead of modeling the full image manifold, PixelGen introduces two complementary perceptual losses to guide diffusion model towards learning a more meaningful **perceptual manifold**. An LPIPS loss facilitates learning better local patterns, while a DINO-based perceptual loss strengthens global semantics. With perceptual supervision, PixelGen surpasses strong latent diffusion baselines. It achieves an FID of 5.11 on ImageNet-256 without classifier-free guidance using only 80 training epochs, and demonstrates favorable scaling performance on large-scale text-to-image generation with a GenEval score of 0.79.
+
++ Figure 1: Visualization of 512x512 images generated by our PixelGen. +
+
+
+
+- We achieve **5.11 FID** on ImageNet256x256 without CFG at 80 epochs, surpassing REPA's 5.90 FID at 800 epochs.
+- We achieve **1.83 FID** on ImageNet256x256 with CFG 160 epochs, competetive with latent diffusion models.
+- We achieve **0.79 overall score** on GenEval Benchmark with PixelGen-XXL/16.
+- **If you like our project, please kindly give us a star ⭐ on GitHub.** We hope to collaborate with you on building better pixel diffusion models, specifically looking at better samplers, CFG strategies, architectures, and refined loss design. Please feel free to reach out if you'd like to discuss ideas.
+
+## +
Illustration of Perceptual Manifold
+
+
+
+ 
+ Illustration of different manifolds within the pixel space. The image manifold is a large manifold containing both perceptually significant information and imperceptible signals. The perceptual manifold contains perceptually important signals, providing a better target for pixel space diffusion. P-DINO and LPIPS are the two complementary perceptual supervision utilized in PixelGen. + +
+
+## 🧩 Visualizations
++ Effectiveness of the perceptual losses in PixelGen.
++ Illustration of different manifolds within the pixel space. The image manifold is a large manifold containing both perceptually significant information and imperceptible signals. The perceptual manifold contains perceptually important signals, providing a better target for pixel space diffusion. P-DINO and LPIPS are the two complementary perceptual supervision utilized in PixelGen. + +
+ 
+
+
++ Visualization of more images generated by our text-to-image PixelGen.
++
+ 
+
+
++ Visualization of 256*256 images generated by our class-to-image PixelGen.
++
+ 
+
+
+## 🎉 Checkpoints
+
+| Dataset | Epoch | Model | Params | FID | HuggingFace |
+|---------------|-------|---------------|--------|-------|---------------------------------------|
+| ImageNet256 | 80 | PixelGen-XL/16 | 676M | 5.11 (w/o CFG) | [🤗](https://huggingface.co/zehongma/PixelGen/blob/main/PixelGen_XL_80ep.ckpt) |
+| ImageNet256 | 160 | PixelGen-XL/16 | 676M | 1.83 (w/ CFG) | [🤗](https://huggingface.co/zehongma/PixelGen/blob/main/PixelGen_XL_160ep.ckpt) |
+
+| Dataset | Model | Params | GenEval | HuggingFace |
+|---------------|---------------|--------|------|----------------------------------------------------------|
+| Text-to-Image | PixelGen-XXL/16| 1.1B | 0.79 | [🤗](https://huggingface.co/zehongma/PixelGen/blob/main/PixelGen_XXL_T2I.ckpt) |
+## 🔥 Online Demos
+
+We provide online demos for PixelGen-XXL/16(text-to-image) on HuggingFace Spaces.
+
+HF spaces: [https://dd0d187fc54e4b00ee.gradio.live](https://dd0d187fc54e4b00ee.gradio.live)
+
+To host the local gradio demo, run the following command:
+```bash
+# for text-to-image applications
+python app.py --config ./configs_t2i/sft_res512.yaml --ckpt_path=./ckpts/PixelGen_XXL_T2I.ckpt
+```
+
+## 🤖 Usages
+In class-to-image(ImageNet) experiments, We use [ADM evaluation suite](https://github.com/openai/guided-diffusion/tree/main/evaluations) to report FID.
+In text-to-image experiments, we use BLIP3o dataset as training set and utilize GenEval to collect metrics.
+
++ Environments
+```bash
+# for installation (recommend python 3.10)
+pip install -r requirements.txt
+```
+
++ Inference
+```bash
+# for inference without CFG using 80-epoch checkpoint
+python main.py predict -c ./configs_c2i/PixelGen_XL_without_CFG.yaml --ckpt_path=./ckpts/PixelGen_XL_80ep.ckpt
+# for inference with CFG using 160-epoch checkpoint
+python main.py predict -c ./configs_c2i/PixelGen_XL.yaml --ckpt_path=./ckpts/PixelGen_XL_160ep.ckpt
+```
+
++ Train
+```bash
+# for c2i training
+# Please modify the ImageNet1k path in the config file before training.
+python main.py fit -c ./configs_c2i/PixelGen_XL.yaml
+```
+
+```bash
+# multi-node training in lightning style, e.g., 4 nodes
+export MASTER_ADDR={Your Config}
+export MASTER_PORT={Your Config}
+export NODE_RANK={Your Config}
+export NNODES={Your Config}
+export NGPUS_PER_NODE={Your Config}
+python main.py fit -c ./configs_c2i/PixelGen_XL.yaml --trainer.num_nodes=4
+```
+
+```bash
+# for t2i training
+python main.py fit -c ./configs_t2i/pretraining_res256.yaml
+python main.py fit -c ./configs_t2i/pretraining_res512.yaml --ckpt_path=./ckpts/pretrain256.ckpt
+python main.py fit -c ./configs_t2i/sft_res512.yaml --ckpt_path=./ckpts/pretrain512.ckpt
+```
+
+## 💐 Acknowledgement
+
+This repository is built based on [DeCo](https://github.com/Zehong-Ma/DeCo). Thanks for their contributions and [Shuai Wang](https://github.com/WANGSSSSSSS)'s support!
+
+## 📖 Citation
+If you find PixelGen is useful in your research or applications, please consider giving us a star ⭐ and citing it by the following BibTeX entry.
+
+```
+@article{ma2026pixelgen,
+ title={PixelGen: Pixel Diffusion Beats Latent Diffusion with Perceptual Loss},
+ author={Zehong Ma and Ruihan Xu and Shiliang Zhang},
+ year={2026},
+ eprint={2602.02493},
+ archivePrefix={arXiv},
+ primaryClass={cs.CV},
+ url={https://arxiv.org/abs/2602.02493},
+}
+```
diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e8d8f32fcfa2487a78cb49e7ed0b2b333a2bea7
--- /dev/null
+++ b/app.py
@@ -0,0 +1,208 @@
+# vae:
+# class_path: src.models.vae.LatentVAE
+# init_args:
+# precompute: true
+# weight_path: /mnt/bn/wangshuai6/models/sd-vae-ft-ema/
+# denoiser:
+# class_path: src.models.denoiser.decoupled_improved_dit.DDT
+# init_args:
+# in_channels: 4
+# patch_size: 2
+# num_groups: 16
+# hidden_size: &hidden_dim 1152
+# num_blocks: 28
+# num_encoder_blocks: 22
+# num_classes: 1000
+# conditioner:
+# class_path: src.models.conditioner.LabelConditioner
+# init_args:
+# null_class: 1000
+# diffusion_sampler:
+# class_path: src.diffusion.stateful_flow_matching.sampling.EulerSampler
+# init_args:
+# num_steps: 250
+# guidance: 3.0
+# state_refresh_rate: 1
+# guidance_interval_min: 0.3
+# guidance_interval_max: 1.0
+# timeshift: 1.0
+# last_step: 0.04
+# scheduler: *scheduler
+# w_scheduler: src.diffusion.stateful_flow_matching.scheduling.LinearScheduler
+# guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+# step_fn: src.diffusion.stateful_flow_matching.sampling.ode_step_fn
+import random
+import os
+import torch
+import argparse
+from omegaconf import OmegaConf
+from src.models.autoencoder.base import fp2uint8
+from src.diffusion.base.guidance import simple_guidance_fn
+from src.diffusion.flow_matching.adam_sampling import AdamLMSamplerJiT
+from src.diffusion.flow_matching.scheduling import LinearScheduler
+from PIL import Image
+import gradio as gr
+import tempfile
+from huggingface_hub import snapshot_download
+
+
+def instantiate_class(config):
+ kwargs = config.get("init_args", {})
+ class_module, class_name = config["class_path"].rsplit(".", 1)
+ module = __import__(class_module, fromlist=[class_name])
+ args_class = getattr(module, class_name)
+ return args_class(**kwargs)
+
+def load_model(weight_dict, denoiser):
+ prefix = "ema_denoiser."
+ for k, v in denoiser.state_dict().items():
+ try:
+ v.copy_(weight_dict["state_dict"][prefix + k])
+ except:
+ print(f"Failed to copy {prefix + k} to denoiser weight")
+ return denoiser
+
+
+class Pipeline:
+ def __init__(self, vae, denoiser, conditioner, resolution):
+ self.vae = vae.cuda()
+ self.denoiser = denoiser.cuda()
+ self.conditioner = conditioner.cuda()
+ self.conditioner.compile()
+ self.resolution = resolution
+ self.tmp_dir = tempfile.TemporaryDirectory(prefix="traj_gifs_")
+ # self.denoiser.compile()
+
+ def __del__(self):
+ self.tmp_dir.cleanup()
+
+ @torch.no_grad()
+ @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+ def __call__(self, y, neg_prompt, num_images, seed, image_height, image_width, num_steps, guidance, timeshift, order):
+ diffusion_sampler = AdamLMSamplerJiT(
+ order=order,
+ scheduler=LinearScheduler(),
+ guidance_fn=simple_guidance_fn,
+ num_steps=num_steps,
+ guidance=guidance,
+ timeshift=timeshift
+ )
+ generator = torch.Generator(device="cpu").manual_seed(seed)
+ image_height = image_height // 32 * 32
+ image_width = image_width // 32 * 32
+ self.denoiser.decoder_patch_scaling_h = image_height / 512
+ self.denoiser.decoder_patch_scaling_w = image_width / 512
+ xT = torch.randn((num_images, 3, image_height, image_width), device="cpu", dtype=torch.float32,
+ generator=generator)
+ xT = xT.to("cuda")
+ with torch.no_grad():
+ condition, uncondition = conditioner([y,]*num_images, {"negative_prompt": neg_prompt})
+
+
+ # Sample images:
+ samples, trajs = diffusion_sampler(denoiser, xT, condition, uncondition, return_x_trajs=True)
+
+ def decode_images(samples):
+ samples = vae.decode(samples)
+ samples = fp2uint8(samples)
+ samples = samples.permute(0, 2, 3, 1).cpu().numpy()
+ images = []
+ for i in range(len(samples)):
+ image = Image.fromarray(samples[i])
+ images.append(image)
+ return images
+
+ def decode_trajs(trajs):
+ cat_trajs = torch.stack(trajs, dim=0).permute(1, 0, 2, 3, 4)
+ animations = []
+ for i in range(cat_trajs.shape[0]):
+ frames = decode_images(
+ cat_trajs[i]
+ )
+ # 生成唯一文件名（结合seed和样本索引，避免冲突）
+ gif_filename = f"{random.randint(0, 100000)}.gif"
+ gif_path = os.path.join(self.tmp_dir.name, gif_filename)
+ frames[0].save(
+ gif_path,
+ format="GIF",
+ append_images=frames[1:],
+ save_all=True,
+ duration=200,
+ loop=0
+ )
+ animations.append(gif_path)
+ return animations
+
+ images = decode_images(samples)
+ animations = decode_trajs(trajs)
+
+ return images, animations
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--config", type=str, default="configs_t2i/sft_res512.yaml")
+ parser.add_argument("--resolution", type=int, default=512)
+ parser.add_argument("--model_id", type=str, default="MCG-NJU/PixNerd-XXL-P16-T2I")
+ parser.add_argument("--ckpt_path", type=str, default="models")
+
+ args = parser.parse_args()
+ if not os.path.exists(args.ckpt_path):
+ snapshot_download(repo_id=args.model_id, local_dir=args.ckpt_path)
+ ckpt_path = os.path.join(args.ckpt_path, "model.ckpt")
+ else:
+ ckpt_path = args.ckpt_path
+
+ config = OmegaConf.load(args.config)
+ vae_config = config.model.vae
+ denoiser_config = config.model.denoiser
+ conditioner_config = config.model.conditioner
+
+ vae = instantiate_class(vae_config)
+ denoiser = instantiate_class(denoiser_config)
+ conditioner = instantiate_class(conditioner_config)
+
+
+ ckpt = torch.load(ckpt_path, map_location="cpu")
+ denoiser = load_model(ckpt, denoiser)
+ denoiser = denoiser.cuda()
+ vae = vae.cuda()
+ denoiser.eval()
+
+
+ pipeline = Pipeline(vae, denoiser, conditioner, args.resolution)
+
+ with gr.Blocks() as demo:
+ # gr.Markdown(f"config:{args.config}\n\n ckpt_path:{args.ckpt_path}")
+ with gr.Row():
+ with gr.Column(scale=1):
+ num_steps = gr.Slider(minimum=1, maximum=100, step=1, label="num steps", value=25)
+ guidance = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, label="CFG", value=4.0)
+ image_height = gr.Slider(minimum=128, maximum=1024, step=32, label="image height", value=512)
+ image_width = gr.Slider(minimum=128, maximum=1024, step=32, label="image width", value=512)
+ num_images = gr.Slider(minimum=1, maximum=4, step=1, label="num images", value=4)
+ label = gr.Textbox(label="positive prompt", value="A beautiful woman.")
+ neg_label = gr.Textbox(label="negative prompt", value="Unrealistic, JPEG artifacts.")
+ seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="seed", value=0)
+ timeshift = gr.Slider(minimum=0.1, maximum=5.0, step=0.1, label="timeshift", value=3.0)
+ order = gr.Slider(minimum=1, maximum=4, step=1, label="order", value=2)
+ with gr.Column(scale=2):
+ btn = gr.Button("Generate")
+ output_sample = gr.Gallery(label="Images", columns=2, rows=2)
+ with gr.Column(scale=2):
+ output_trajs = gr.Gallery(label="Trajs of Diffusion", columns=2, rows=2)
+
+ btn.click(fn=pipeline,
+ inputs=[
+ label,
+ neg_label,
+ num_images,
+ seed,
+ image_height,
+ image_width,
+ num_steps,
+ guidance,
+ timeshift,
+ order
+ ], outputs=[output_sample, output_trajs])
+ # demo.launch(server_name="0.0.0.0", server_port=23231)
+ demo.launch(share=True, server_name="0.0.0.0", server_port=23231)
\ No newline at end of file
diff --git a/ck/PixelGen_XXL_T2I.ckpt b/ck/PixelGen_XXL_T2I.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..c809471c2826c830276c5101482b52bf8f9566d2
--- /dev/null
+++ b/ck/PixelGen_XXL_T2I.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b3e10277aa17dd033dc34b1f9e35756e40ec1cbee0d11cbdfb376becc78e55c
+size 4516959517
diff --git a/config.yaml b/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..3e59de9d5cd3738badd9d59ce9486b7adeee20a4
--- /dev/null
+++ b/config.yaml
@@ -0,0 +1,290 @@
+# lightning.pytorch==2.6.1
+seed_everything: 0
+torch_hub_dir: ~/.cache/torch/hub
+huggingface_cache_dir: null
+trainer:
+ accelerator: auto
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ name: sft_pix512_gan_v8
+ save_dir: .
+ version: null
+ offline: false
+ dir: null
+ id: null
+ anonymous: null
+ project: universal_flow_JiT_t2i
+ log_model: false
+ experiment: null
+ prefix: ''
+ checkpoint_name: null
+ add_file_policy: mutable
+ entity: null
+ notes: null
+ tags: null
+ config: null
+ config_exclude_keys: null
+ config_include_keys: null
+ allow_val_change: null
+ group: null
+ job_type: null
+ mode: null
+ force: null
+ reinit: null
+ resume: null
+ resume_from: null
+ fork_from: null
+ save_code: null
+ tensorboard: null
+ sync_tensorboard: null
+ monitor_gym: null
+ settings: null
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ dirpath: null
+ filename: null
+ monitor: null
+ verbose: false
+ save_last: true
+ save_top_k: -1
+ save_on_exception: false
+ save_weights_only: false
+ mode: min
+ auto_insert_metric_name: true
+ every_n_train_steps: 5000
+ train_time_interval: null
+ every_n_epochs: null
+ save_on_train_epoch_end: null
+ enable_version_counter: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ save_compressed: false
+ fast_dev_run: false
+ max_epochs: null
+ min_epochs: null
+ max_steps: 200000
+ min_steps: null
+ max_time: null
+ limit_train_batches: null
+ limit_val_batches: null
+ limit_test_batches: null
+ limit_predict_batches: null
+ overfit_batches: 0.0
+ val_check_interval: 10000
+ check_val_every_n_epoch: null
+ num_sanity_val_steps: 2
+ log_every_n_steps: 50
+ enable_checkpointing: null
+ enable_progress_bar: null
+ enable_model_summary: null
+ accumulate_grad_batches: 1
+ gradient_clip_val: null
+ gradient_clip_algorithm: null
+ deterministic: null
+ benchmark: null
+ inference_mode: true
+ use_distributed_sampler: false
+ profiler: null
+ detect_anomaly: false
+ barebones: false
+ plugins:
+ - class_path: src.plugins.bd_env.BDEnvironment
+ sync_batchnorm: false
+ reload_dataloaders_every_n_epochs: 0
+ default_root_dir: ./universal_pix_t2i_workdirs
+ enable_autolog_hparams: true
+ model_registry: null
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ init_args:
+ scale: 1.0
+ shift: 0.0
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: 2048
+ max_length: 128
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ input_size: 512
+ in_channels: 3
+ num_groups: 24
+ hidden_size: 1536
+ num_blocks: 16
+ num_text_blocks: 4
+ patch_size: 16
+ txt_embed_dim: 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ weight_path: null
+ load_ema: false
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_DMD2_GAN.REPADMD2GANTrainer
+ init_args:
+ scheduler:
+ class_path: src.diffusion.flow_matching.scheduling.LinearScheduler
+ loss_weight_fn: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating.constant
+ feat_loss_weight: 0.5
+ lognorm_t: true
+ timeshift: 1.0
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ base_patch_size: 16
+ align_layer: 8
+ proj_denoiser_dim: 1536
+ proj_hidden_dim: 1536
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ num_classes: 1000
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ noise_scale: 1.0
+ patch_size: 16
+ percept_ratio: 1.0
+ percept_adaptive_step: 0
+ use_gan_loss: true
+ gan_loss_weight: 0.003
+ d_fm_loss_weight: 1.0
+ gan_loss_weight_final: 0.003
+ decay_start_step: 0
+ decay_end_step: 0
+ lr_final_multiplier: 1.0
+ decay_schedule: constant
+ diffusion_gan: true
+ diffusion_gan_max_timestep: 0.4
+ gan_start_step: 0
+ gan_timestep_min: 0.6
+ gan_timestep_max: 1.0
+ null_condition_p: 0.1
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ order: 2
+ timeshift: 3.0
+ guidance_interval_min: 0.0
+ guidance_interval_max: 1.0
+ lms_transform_fn: src.diffusion.flow_matching.adam_sampling.nop
+ last_step: null
+ step_fn: src.diffusion.flow_matching.adam_sampling.ode_step_fn
+ scheduler:
+ class_path: src.diffusion.flow_matching.scheduling.LinearScheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ num_steps: 25
+ guidance: 4.0
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ every_n_steps: 1
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-05
+ betas:
+ - 0.9
+ - 0.999
+ eps: 1.0e-08
+ weight_decay: 0.0
+ amsgrad: false
+ maximize: false
+ foreach: null
+ capturable: false
+ differentiable: false
+ fused: null
+ lr_scheduler: null
+ eval_original_model: false
+ discriminator:
+ class_path: src.models.discriminator.DenoiserBackboneDiscriminator
+ init_args:
+ backbone:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ input_size: 512
+ in_channels: 3
+ num_groups: 24
+ hidden_size: 1536
+ num_blocks: 16
+ num_text_blocks: 4
+ patch_size: 16
+ txt_embed_dim: 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ weight_path: null
+ load_ema: false
+ hidden_size: 1536
+ head_hidden_size: 1536
+ head_depth: 2
+ train_backbone: true
+ backbone_lr_scale: 1.0
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-05
+ betas:
+ - 0.0
+ - 0.99
+ eps: 1.0e-08
+ weight_decay: 0.0
+ amsgrad: false
+ maximize: false
+ foreach: null
+ capturable: false
+ differentiable: false
+ fused: null
+ d_steps_per_g: 1
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ cache_path: null
+ require_text: true
+ eval_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ limit: 50
+ seed_offset: 0
+ pred_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ limit: 50
+ seed_offset: 0
+ train_batch_size: 8
+ train_num_workers: 8
+ train_prefetch_factor: 8
+ eval_batch_size: 32
+ eval_num_workers: 4
+ pred_batch_size: 8
+ pred_num_workers: 1
+tags:
+ exp: sft_pix512_gan_v8
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+weights_only: null
diff --git a/configs_c2i/JiT_Large_Baseline.yaml b/configs_c2i/JiT_Large_Baseline.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..44a3216ae5e77b4948e566f125bd81b64c05fcf4
--- /dev/null
+++ b/configs_c2i/JiT_Large_Baseline.yaml
@@ -0,0 +1,134 @@
+# lightning.pytorch==2.4.0
+seed_everything: true
+tags:
+ exp: &exp JiT_REPA_256_baseline
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_workdirs
+ accelerator: auto
+ strategy: auto
+ devices: auto
+ num_nodes: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_bs256
+ name: *exp
+ num_sanity_val_steps: 0
+ max_steps: 200000
+ val_check_interval: 100000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 100000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val_ode50_cfg2.0
+ save_compressed: true
+ # - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+ # init_args:
+ # logging_interval: 'step'
+ # log_momentum: False
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ init_args:
+ scale: 1.0
+ denoiser:
+ class_path: src.models.transformer.JiT.JiT
+ init_args:
+ input_size: 256
+ patch_size: 16
+ in_channels: 3
+ hidden_size: &hidden_dim 1024
+ depth: 24
+ num_heads: 16
+ mlp_ratio: 4.0
+ attn_drop: 0.0
+ proj_drop: 0.0
+ num_classes: &num_classes 1000
+ bottleneck_dim: 128
+ in_context_len: 32
+ in_context_start: 8
+ conditioner:
+ class_path: src.models.conditioner.class_label.LabelConditioner
+ init_args:
+ num_classes: *num_classes
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT.REPATrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.sampling.EulerSamplerJiT
+ init_args:
+ num_steps: 50
+ guidance: 1.0
+ guidance_interval_min: 0.1
+ guidance_interval_max: 1.0
+ scheduler: *scheduler
+ w_scheduler: src.diffusion.flow_matching.scheduling.LinearScheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ step_fn: src.diffusion.flow_matching.sampling.ode_step_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1e-4
+ weight_decay: 0.0
+ # lr_scheduler:
+ # class_path: src.utils.lr_scheduler.ConstantWithWarmup
+ # init_args:
+ # num_warmup_steps: 5005
+data:
+ train_dataset:
+ class_path: src.data.dataset.imagenet.PixImageNet
+ init_args:
+ root: /data/datasets/ImageNet-1K/train
+ resolution: 256
+ eval_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: 1000
+ max_num_instances: 50000
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ pred_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: *num_classes
+ max_num_instances: 50000
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ train_batch_size: 32
+ train_num_workers: 4
+ pred_batch_size: 32
+ pred_num_workers: 1
diff --git a/configs_c2i/PixelGen_Large.yaml b/configs_c2i/PixelGen_Large.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..1f14966cd5b243f5ca8870f86c114292898675be
--- /dev/null
+++ b/configs_c2i/PixelGen_Large.yaml
@@ -0,0 +1,137 @@
+# lightning.pytorch==2.4.0
+seed_everything: true
+tags:
+ exp: &exp PixelGen_Large
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_workdirs
+ accelerator: auto
+ strategy: auto
+ devices: auto
+ num_nodes: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_bs256
+ name: *exp
+ num_sanity_val_steps: 0
+ max_steps: 200000
+ val_check_interval: 100000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 100000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val_ode50_cfg1.0
+ save_compressed: true
+ # - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+ # init_args:
+ # logging_interval: 'step'
+ # log_momentum: False
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ init_args:
+ scale: 1.0
+ denoiser:
+ class_path: src.models.transformer.JiT.JiT
+ init_args:
+ input_size: 256
+ patch_size: 16
+ in_channels: 3
+ hidden_size: &hidden_dim 1024
+ depth: 24
+ num_heads: 16
+ mlp_ratio: 4.0
+ attn_drop: 0.0
+ proj_drop: 0.0
+ num_classes: &num_classes 1000
+ bottleneck_dim: 128
+ in_context_len: 32
+ in_context_start: 8
+ conditioner:
+ class_path: src.models.conditioner.class_label.LabelConditioner
+ init_args:
+ num_classes: *num_classes
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating.REPATrainer # see other trainers for ablation experiments.
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1 # LPIPS loss, 0.0 to disable
+ dino_weight: 0.01 # P-DINO loss, 0.0 to disable
+ percept_t_threshold: 0.3 # Noise-Gating, 0.0 to disable
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.sampling.EulerSamplerJiT
+ init_args:
+ num_steps: 50
+ guidance: 1.0
+ guidance_interval_min: 0.1
+ guidance_interval_max: 1.0
+ scheduler: *scheduler
+ w_scheduler: src.diffusion.flow_matching.scheduling.LinearScheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ step_fn: src.diffusion.flow_matching.sampling.ode_step_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1e-4
+ weight_decay: 0.0
+ # lr_scheduler:
+ # class_path: src.utils.lr_scheduler.ConstantWithWarmup
+ # init_args:
+ # num_warmup_steps: 5005
+data:
+ train_dataset:
+ class_path: src.data.dataset.imagenet.PixImageNet
+ init_args:
+ root: /data/datasets/ImageNet-1K/train
+ resolution: 256
+ eval_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: 1000
+ max_num_instances: 50000
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ pred_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: *num_classes
+ max_num_instances: 50000
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ train_batch_size: 32
+ train_num_workers: 4
+ pred_batch_size: 32
+ pred_num_workers: 1
diff --git a/configs_c2i/PixelGen_XL.yaml b/configs_c2i/PixelGen_XL.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..bca33ec597b99d1af00c40d7e9527fa6e4d6e3b0
--- /dev/null
+++ b/configs_c2i/PixelGen_XL.yaml
@@ -0,0 +1,141 @@
+# lightning.pytorch==2.4.0
+seed_everything: 1234
+tags:
+ exp: &exp PixelGen_XL
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_workdirs
+ accelerator: auto
+ strategy: auto
+ devices: auto
+ num_nodes: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_XL_bs256
+ name: *exp
+ num_sanity_val_steps: 0
+ max_steps: 500000 # 800k for 160 epochs, 4000k for 800 epochs
+ val_check_interval: 100000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 100000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val_ode50_cfg2.25
+ save_compressed: true
+ # - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+ # init_args:
+ # logging_interval: 'step'
+ # log_momentum: False
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ init_args:
+ scale: 1.0
+ denoiser:
+ class_path: src.models.transformer.JiT.JiT
+ init_args:
+ input_size: 256
+ patch_size: 16
+ in_channels: 3
+ hidden_size: &hidden_dim 1152
+ depth: 28
+ num_heads: 16
+ mlp_ratio: 4.0
+ attn_drop: 0.0
+ proj_drop: 0.1 # please enable proj_drop when training a DiT-XL with lots of parameters
+ num_classes: &num_classes 1000
+ use_bottleneck: false
+ in_context_len: 32
+ in_context_start: 8
+ conditioner:
+ class_path: src.models.conditioner.class_label.LabelConditioner
+ init_args:
+ num_classes: *num_classes
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating.REPATrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ null_condition_p: 0.1
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.sampling.HeunSamplerJiT
+ init_args:
+ exact_henu: true
+ num_steps: 50
+ guidance: 2.25 # Guidance scale is set to 2.25 for 160 epochs. Please set to 1.0 for the experiments without CFG.
+ timeshift: 2.0 # For XL version, the timesift is set to 2.0 during all inference.
+ guidance_interval_min: 0.1
+ guidance_interval_max: 0.9
+ scheduler: *scheduler
+ w_scheduler: src.diffusion.flow_matching.scheduling.LinearScheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ step_fn: src.diffusion.flow_matching.sampling.ode_step_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1e-4
+ weight_decay: 0.0
+ # lr_scheduler:
+ # class_path: src.utils.lr_scheduler.ConstantWithWarmup
+ # init_args:
+ # num_warmup_steps: 5005
+data:
+ train_dataset:
+ class_path: src.data.dataset.imagenet.PixImageNet
+ init_args:
+ root: /data/datasets/ImageNet-1K/train # set your data dir
+ resolution: 256
+ eval_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: 1000
+ max_num_instances: 50000
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ pred_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: *num_classes
+ max_num_instances: 50000
+ noise_scale: 1.0
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ train_batch_size: 32
+ train_num_workers: 4
+ pred_batch_size: 32
+ pred_num_workers: 1
diff --git a/configs_c2i/PixelGen_XL_without_CFG.yaml b/configs_c2i/PixelGen_XL_without_CFG.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..1daaddb5ce315d93020521ce952626d47c413567
--- /dev/null
+++ b/configs_c2i/PixelGen_XL_without_CFG.yaml
@@ -0,0 +1,141 @@
+# lightning.pytorch==2.4.0
+seed_everything: true
+tags:
+ exp: &exp PixelGen_XL
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_workdirs
+ accelerator: auto
+ strategy: auto
+ devices: auto
+ num_nodes: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_XL_bs256
+ name: *exp
+ num_sanity_val_steps: 0
+ max_steps: 800000 # 800k for 160 epochs, 4000k for 800 epochs
+ val_check_interval: 100000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 100000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val_ode50_cfg1.0
+ save_compressed: true
+ # - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+ # init_args:
+ # logging_interval: 'step'
+ # log_momentum: False
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ init_args:
+ scale: 1.0
+ denoiser:
+ class_path: src.models.transformer.JiT.JiT
+ init_args:
+ input_size: 256
+ patch_size: 16
+ in_channels: 3
+ hidden_size: &hidden_dim 1152
+ depth: 28
+ num_heads: 16
+ mlp_ratio: 4.0
+ attn_drop: 0.0
+ proj_drop: 0.1 # please enable proj_drop when training a DiT-XL with lots of parameters
+ num_classes: &num_classes 1000
+ use_bottleneck: false
+ in_context_len: 32
+ in_context_start: 8
+ conditioner:
+ class_path: src.models.conditioner.class_label.LabelConditioner
+ init_args:
+ num_classes: *num_classes
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating.REPATrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ null_condition_p: 0.1
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.sampling.HeunSamplerJiT
+ init_args:
+ exact_henu: true
+ num_steps: 50
+ guidance: 1.0 # Guidance scale is set to 1.0 for inferring with 80 epochs ckpt without CFG.
+ timeshift: 2.0 # For XL version, the timesift is set to 2.0 during all inference.
+ guidance_interval_min: 0.1
+ guidance_interval_max: 0.9
+ scheduler: *scheduler
+ w_scheduler: src.diffusion.flow_matching.scheduling.LinearScheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ step_fn: src.diffusion.flow_matching.sampling.ode_step_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1e-4
+ weight_decay: 0.0
+ # lr_scheduler:
+ # class_path: src.utils.lr_scheduler.ConstantWithWarmup
+ # init_args:
+ # num_warmup_steps: 5005
+data:
+ train_dataset:
+ class_path: src.data.dataset.imagenet.PixImageNet
+ init_args:
+ root: /data/datasets/ImageNet-1K/train # set your data dir
+ resolution: 256
+ eval_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: 1000
+ max_num_instances: 50000
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ pred_dataset:
+ class_path: src.data.dataset.randn.ClassLabelRandomNDataset
+ init_args:
+ num_classes: *num_classes
+ max_num_instances: 50000
+ noise_scale: 1.0
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ train_batch_size: 32
+ train_num_workers: 4
+ pred_batch_size: 32
+ pred_num_workers: 1
diff --git a/configs_t2i/pretraining_res256.yaml b/configs_t2i/pretraining_res256.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..8f2103b517a7f81ea2c3bc3e1ef5f78052beae92
--- /dev/null
+++ b/configs_t2i/pretraining_res256.yaml
@@ -0,0 +1,134 @@
+# lightning.pytorch==2.4.0
+seed_everything: true
+tags:
+ exp: &exp pretraining_pix256
+torch_hub_dir: ~/.cache/torch/hub
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 3
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ gradient_clip_val: 1.0
+ gradient_clip_algorithm: norm
+ val_check_interval: 20000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 20000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating.REPATrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.sampling.EulerSamplerJiT
+ init_args:
+ num_steps: 100
+ guidance: 4.0
+ scheduler: *scheduler
+ w_scheduler: src.diffusion.flow_matching.scheduling.LinearScheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ step_fn: src.diffusion.flow_matching.sampling.ode_step_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 2e-4
+ betas:
+ - 0.9
+ - 0.95
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.blip3o_dataset.WebDatasetPackedDataset
+ init_args:
+ urls: [/data/datasets/BLIP-3o/BLIP3o-Pretrain-Long-Caption, /data/datasets/BLIP-3o/BLIP3o-Pretrain-Short-Caption, /data/datasets/BLIP-3o/BLIP3o-Pretrain-JourneyDB] # [/data/datasets/BLIP-3o/BLIP3o-Pretrain-JourneyDB,]
+ resolution: 256
+ random_crop: false
+ shuffle_buffer: 20000
+ sample_shuffle: true
+ repeat: true
+ eval_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ pred_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ # meta_json_path: ./evaluations/geneval/evaluation_metadata_rephrased.jsonl
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 256
+ - 256
+ train_batch_size: 64 # 96
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
\ No newline at end of file
diff --git a/configs_t2i/pretraining_res512.yaml b/configs_t2i/pretraining_res512.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..cd77ecdaa855ebf412c6513e499318214b10088b
--- /dev/null
+++ b/configs_t2i/pretraining_res512.yaml
@@ -0,0 +1,136 @@
+# lightning.pytorch==2.4.0
+seed_everything: true
+tags:
+ exp: &exp pretraining_pix512
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /data/github_project/DeCo/universal_pix_t2i_workdirs/exp_pretraining_pix256/epoch=0-step=200000.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 4
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 300000
+ gradient_clip_val: 1.0
+ gradient_clip_algorithm: norm
+ val_check_interval: 40000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 10000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating.REPATrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.sampling.EulerSamplerJiT
+ init_args:
+ num_steps: 100
+ guidance: 4.0
+ scheduler: *scheduler
+ w_scheduler: src.diffusion.flow_matching.scheduling.LinearScheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ step_fn: src.diffusion.flow_matching.sampling.ode_step_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1e-4
+ betas:
+ - 0.9
+ - 0.95
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.blip3o_dataset.WebDatasetPackedDataset
+ init_args:
+ urls: [/data/datasets/BLIP-3o/BLIP3o-Pretrain-Long-Caption, /data/datasets/BLIP-3o/BLIP3o-Pretrain-Short-Caption, /data/datasets/BLIP-3o/BLIP3o-Pretrain-JourneyDB] # [/data/datasets/BLIP-3o/BLIP3o-Pretrain-JourneyDB,]
+ resolution: 512
+ random_crop: false
+ shuffle_buffer: 20000
+ sample_shuffle: true
+ repeat: true
+ eval_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ # meta_json_path: ./evaluations/geneval/evaluation_metadata_rephrased.jsonl
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 16
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
\ No newline at end of file
diff --git a/configs_t2i/sft_res512.yaml b/configs_t2i/sft_res512.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..af23ffa90e86ef8669ff28f29d93def838a1d1f2
--- /dev/null
+++ b/configs_t2i/sft_res512.yaml
@@ -0,0 +1,148 @@
+# lightning.pytorch==2.4.0
+seed_everything: true
+tags:
+ exp: &exp sftpix512ema
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /data/github_project/DeCo/universal_pix_t2i_workdirs/exp_pretraining_pix512/epoch=0-step=300000.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 4
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 360000
+ gradient_clip_val: 1.0
+ gradient_clip_algorithm: norm
+ val_check_interval: 40000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 10000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating.REPATrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0 # 3.0 for geneval
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /data/datasets/BLIP-3o/BLIP3o-60k
+ resolution: 512
+ # train_dataset:
+ # class_path: src.data.dataset.blip3o_dataset.WebDatasetPackedDataset
+ # init_args:
+ # urls: [/data/datasets/BLIP-3o/BLIP3o-60k] # [/data/datasets/BLIP-3o/BLIP3o-Pretrain-JourneyDB,]
+ # resolution: 512
+ # random_crop: false
+ # shuffle_buffer: 2000
+ # sample_shuffle: true
+ # repeat: true
+ eval_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ # class_path: src.data.dataset.dpg.DPGDataset
+ # init_args:
+ # prompt_path: ./evaluations/dpg/prompts
+ # num_samples_per_instance: 4
+ # latent_shape:
+ # - 3
+ # - 512
+ # - 512
+ train_batch_size: 16
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
\ No newline at end of file
diff --git a/configs_t2i/sft_res512_gan.yaml b/configs_t2i/sft_res512_gan.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..48c9a8aa6a536531abab678361d9b479d50c9a57
--- /dev/null
+++ b/configs_t2i/sft_res512_gan.yaml
@@ -0,0 +1,194 @@
+# lightning.pytorch==2.4.0
+# ---------------------------------------------------------------------------
+# PixelGen-XXL T2I SFT @ 512 + StyleGAN-T-style projected GAN loss.
+#
+# What is new vs. configs_t2i/sft_res512.yaml:
+# - diffusion_trainer is the GAN-aware REPAGANTrainer (REPA + LPIPS + DINO
+# + projected hinge GAN loss).
+# - A `discriminator` block instantiates a ProjectedDiscriminator: a *frozen*
+# DINOv2 ViT-B/14 backbone topped with a few trainable per-layer CNN heads
+# plus a trainable text-projection MLP.
+# - Manual optimization is auto-enabled inside LightningModel because a
+# discriminator is supplied; gradient clipping must therefore be removed
+# from the Trainer block and is performed inside training_step instead.
+# ---------------------------------------------------------------------------
+seed_everything: true
+tags:
+ exp: &exp sft_pix512_gan_v2
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ # GAN training does two manual_backward calls per training_step (G then D).
+ # DDP would otherwise complain that some params didn't participate in each
+ # backward, so we use the "find_unused_parameters=True" DDP variant.
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ # NOTE: manual optimization disables Lightning's accumulate_grad_batches.
+ # Keep it at 1; raise per-GPU batch size or device count instead.
+ accumulate_grad_batches: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ # gradient_clip_val / gradient_clip_algorithm intentionally omitted:
+ # the model performs per-optimizer clipping itself via g_grad_clip / d_grad_clip.
+ val_check_interval: 20000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 5000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_GAN.REPAGANTrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ # We are fine-tuning a pretrained model, not training from scratch, so we
+ # immediately enable the gradient-norm based adaptive balancing between
+ # the FM loss and the perceptual losses (LPIPS + DINO). In the original
+ # schedule this only kicks in after 400k steps.
+ percept_adaptive_step: 0
+ # ----- GAN-specific -----
+ # Lower the GAN weight: with a frozen-DINO projected discriminator the GAN
+ # pulls the generator toward sharp / posterized "DINO-friendly" textures
+ # quite aggressively, so a smaller weight is safer for fine-tuning.
+ gan_weight: 0.03
+ # t convention: t=1 is clean image, t=0 is pure noise. Only let the GAN
+ # supervise lower-noise predictions (later in the trajectory) where
+ # pred_img is already a reasonable image. Roughly the last ~40% of the
+ # trajectory will participate in GAN.
+ gan_t_threshold: 0.6
+ # Already a fine-tune of a pretrained text-to-image model, no warmup
+ # needed before the GAN starts contributing.
+ gan_warmup_steps: 0
+ # ---- NEW: projected discriminator ----
+ discriminator:
+ class_path: src.models.discriminator.ProjectedDiscriminator
+ init_args:
+ dino_encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ dino_layers: [2, 5, 8, 11]
+ embed_dim: 768
+ text_dim: *txt_embed_dim
+ cmap_dim: 64
+ diffaug: true
+ diffaug_policy: "color,translation,cutout"
+ p_crop: 0.5
+ crop_resolution: 224
+ input_resolution: 224
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 2.0e-4
+ betas:
+ - 0.0
+ - 0.99
+ weight_decay: 0.00
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ eval_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 16
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
diff --git a/configs_t2i/sft_res512_gan_v3.yaml b/configs_t2i/sft_res512_gan_v3.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..c6760622ebb2b3076450539b6c0bc31b3e16eeca
--- /dev/null
+++ b/configs_t2i/sft_res512_gan_v3.yaml
@@ -0,0 +1,166 @@
+# lightning.pytorch==2.4.0
+# ---------------------------------------------------------------------------
+# PixelGen-XXL T2I SFT @ 512 + StyleGAN-T-style projected GAN loss.
+#
+# v3 change vs v2:
+# - Lower gan_weight from 0.03 to 0.01 for a more conservative adversarial
+# fine-tuning setup.
+# ---------------------------------------------------------------------------
+seed_everything: true
+tags:
+ exp: &exp sft_pix512_gan_v3
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ val_check_interval: 20000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 5000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_GAN.REPAGANTrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ percept_adaptive_step: 0
+ gan_weight: 0.01
+ gan_t_threshold: 0.6
+ gan_warmup_steps: 0
+ discriminator:
+ class_path: src.models.discriminator.ProjectedDiscriminator
+ init_args:
+ dino_encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ dino_layers: [2, 5, 8, 11]
+ embed_dim: 768
+ text_dim: *txt_embed_dim
+ cmap_dim: 64
+ diffaug: true
+ diffaug_policy: "color,translation,cutout"
+ p_crop: 0.5
+ crop_resolution: 224
+ input_resolution: 224
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 2.0e-4
+ betas:
+ - 0.0
+ - 0.99
+ weight_decay: 0.00
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ eval_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 16
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
diff --git a/configs_t2i/sft_res512_gan_v4.yaml b/configs_t2i/sft_res512_gan_v4.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..92529d59f8629da4a00d633ec351db82140fdf9e
--- /dev/null
+++ b/configs_t2i/sft_res512_gan_v4.yaml
@@ -0,0 +1,179 @@
+# lightning.pytorch==2.4.0
+# ---------------------------------------------------------------------------
+# PixelGen-XXL T2I SFT @ 512 + DMD2-style GAN branch.
+#
+# Built from the original PixelGen SFT config/trainer path:
+# training_repa_JiT_LPIPS_DINO_NoiseGating.REPATrainer
+# and extends it with a discriminator that reuses a trainable copy of the
+# PixelGen DiT denoiser backbone plus a small randomly initialized MLP head.
+# It does not use the v2/v3 ProjectedDINO discriminator.
+# ---------------------------------------------------------------------------
+seed_everything: true
+tags:
+ exp: &exp sft_pix512_gan_v4
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ val_check_interval: 20000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 5000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_DMD2_GAN.REPADMD2GANTrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ percept_adaptive_step: 0
+ use_gan_loss: true
+ gan_loss_weight: 0.01
+ diffusion_gan: true
+ # Add random flow-matching noise before D, DMD2-style. In this codebase
+ # t=1 is clean and t=0 is pure noise, so max 0.4 keeps D inputs in the
+ # same lower-noise region as gan_timestep_min=0.6.
+ diffusion_gan_max_timestep: 0.4
+ gan_start_step: 0
+ gan_timestep_min: 0.6
+ gan_timestep_max: 1.0
+ discriminator:
+ class_path: src.models.discriminator.DenoiserBackboneDiscriminator
+ init_args:
+ backbone:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: *hidden_dim
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: *txt_embed_dim
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ hidden_size: *hidden_dim
+ head_hidden_size: *hidden_dim
+ head_depth: 2
+ train_backbone: false
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 2.0e-4
+ betas:
+ - 0.0
+ - 0.99
+ weight_decay: 0.00
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ eval_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 16
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
diff --git a/configs_t2i/sft_res512_gan_v5.yaml b/configs_t2i/sft_res512_gan_v5.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..52236c9ca51571c1e81296ba4a7278f53a83e29e
--- /dev/null
+++ b/configs_t2i/sft_res512_gan_v5.yaml
@@ -0,0 +1,184 @@
+# lightning.pytorch==2.4.0
+# ---------------------------------------------------------------------------
+# PixelGen-XXL T2I SFT @ 512 + DMD2-style GAN branch.
+#
+# v5 change vs v4:
+# - Keep GAN weight and G/D learning rates unchanged for the first 10k steps.
+# - Cosine decay both GAN weight and G/D learning rates from 10k to 30k.
+# - Hold at 30% strength after 30k because visual quality peaked near 10k and
+# started degrading by 30k in v4.
+# ---------------------------------------------------------------------------
+seed_everything: true
+tags:
+ exp: &exp sft_pix512_gan_v5
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ val_check_interval: 20000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 5000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_DMD2_GAN.REPADMD2GANTrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ percept_adaptive_step: 0
+ use_gan_loss: true
+ gan_loss_weight: 0.01
+ gan_loss_weight_final: 0.003
+ decay_start_step: 10000
+ decay_end_step: 30000
+ lr_final_multiplier: 0.3
+ decay_schedule: cosine
+ diffusion_gan: true
+ # Add random flow-matching noise before D, DMD2-style. In this codebase
+ # t=1 is clean and t=0 is pure noise, so max 0.4 keeps D inputs in the
+ # same lower-noise region as gan_timestep_min=0.6.
+ diffusion_gan_max_timestep: 0.4
+ gan_start_step: 0
+ gan_timestep_min: 0.6
+ gan_timestep_max: 1.0
+ discriminator:
+ class_path: src.models.discriminator.DenoiserBackboneDiscriminator
+ init_args:
+ backbone:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: *hidden_dim
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: *txt_embed_dim
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ hidden_size: *hidden_dim
+ head_hidden_size: *hidden_dim
+ head_depth: 2
+ train_backbone: false
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 2.0e-4
+ betas:
+ - 0.0
+ - 0.99
+ weight_decay: 0.00
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ eval_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.geneval.GenEvalDataset
+ init_args:
+ meta_json_path: ./evaluations/geneval/evaluation_metadata.jsonl
+ num_samples_per_instance: 4
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 16
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
diff --git a/configs_t2i/sft_res512_gan_v6.yaml b/configs_t2i/sft_res512_gan_v6.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..98884d967c2f52a50bbaf320bdfd70eeeca4d121
--- /dev/null
+++ b/configs_t2i/sft_res512_gan_v6.yaml
@@ -0,0 +1,186 @@
+# lightning.pytorch==2.4.0
+# ---------------------------------------------------------------------------
+# PixelGen-XXL T2I SFT @ 512 + DMD2-style GAN branch.
+#
+# v6 change vs v5:
+# - Keep the v5 GAN-weight / G-D learning-rate cosine decay schedule.
+# - Unfreeze the discriminator PixelGen backbone.
+# - Train discriminator head at D lr and backbone at 0.1x D lr.
+# ---------------------------------------------------------------------------
+seed_everything: true
+tags:
+ exp: &exp sft_pix512_gan_v6
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ val_check_interval: 20000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 5000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_DMD2_GAN.REPADMD2GANTrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ percept_adaptive_step: 0
+ use_gan_loss: true
+ gan_loss_weight: 0.01
+ gan_loss_weight_final: 0.003
+ decay_start_step: 10000
+ decay_end_step: 30000
+ lr_final_multiplier: 0.3
+ decay_schedule: cosine
+ diffusion_gan: true
+ # Add random flow-matching noise before D, DMD2-style. In this codebase
+ # t=1 is clean and t=0 is pure noise, so max 0.4 keeps D inputs in the
+ # same lower-noise region as gan_timestep_min=0.6.
+ diffusion_gan_max_timestep: 0.4
+ gan_start_step: 0
+ gan_timestep_min: 0.6
+ gan_timestep_max: 1.0
+ discriminator:
+ class_path: src.models.discriminator.DenoiserBackboneDiscriminator
+ init_args:
+ backbone:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: *hidden_dim
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: *txt_embed_dim
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ hidden_size: *hidden_dim
+ head_hidden_size: *hidden_dim
+ head_depth: 2
+ train_backbone: true
+ backbone_lr_scale: 0.1
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 2.0e-4
+ betas:
+ - 0.0
+ - 0.99
+ weight_decay: 0.00
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ eval_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ limit: 50
+ seed_offset: 0
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ limit: 50
+ seed_offset: 0
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 12
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
diff --git a/configs_t2i/sft_res512_gan_v7.yaml b/configs_t2i/sft_res512_gan_v7.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..1250a30c576a4a5d85eedef2ab0fa0a097d07fc1
--- /dev/null
+++ b/configs_t2i/sft_res512_gan_v7.yaml
@@ -0,0 +1,188 @@
+# lightning.pytorch==2.4.0
+# ---------------------------------------------------------------------------
+# PixelGen-XXL T2I SFT @ 512 + DMD2-style GAN branch.
+#
+# v7 change vs v6:
+# - Keep discriminator backbone trainable.
+# - Train discriminator head/backbone with the same base LR as generator.
+# - Add an FM loss to the discriminator backbone, so D optimizes:
+# D loss = GAN discriminator loss + d_fm_loss_weight * D-backbone FM loss.
+# ---------------------------------------------------------------------------
+seed_everything: true
+tags:
+ exp: &exp sft_pix512_gan_v7
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ val_check_interval: 20000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 5000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_DMD2_GAN.REPADMD2GANTrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ percept_adaptive_step: 0
+ use_gan_loss: true
+ gan_loss_weight: 0.01
+ d_fm_loss_weight: 1.0
+ gan_loss_weight_final: 0.003
+ decay_start_step: 10000
+ decay_end_step: 30000
+ lr_final_multiplier: 0.3
+ decay_schedule: cosine
+ diffusion_gan: true
+ # Add random flow-matching noise before D, DMD2-style. In this codebase
+ # t=1 is clean and t=0 is pure noise, so max 0.4 keeps D inputs in the
+ # same lower-noise region as gan_timestep_min=0.6.
+ diffusion_gan_max_timestep: 0.4
+ gan_start_step: 0
+ gan_timestep_min: 0.6
+ gan_timestep_max: 1.0
+ discriminator:
+ class_path: src.models.discriminator.DenoiserBackboneDiscriminator
+ init_args:
+ backbone:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: *hidden_dim
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: *txt_embed_dim
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ hidden_size: *hidden_dim
+ head_hidden_size: *hidden_dim
+ head_depth: 2
+ train_backbone: true
+ backbone_lr_scale: 1.0
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.0
+ - 0.99
+ weight_decay: 0.00
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ eval_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ limit: 50
+ seed_offset: 0
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ limit: 50
+ seed_offset: 0
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 8
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
diff --git a/configs_t2i/sft_res512_gan_v8.yaml b/configs_t2i/sft_res512_gan_v8.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..480bd33d8eaf087e1d5a80002270343da5457019
--- /dev/null
+++ b/configs_t2i/sft_res512_gan_v8.yaml
@@ -0,0 +1,190 @@
+# lightning.pytorch==2.4.0
+# ---------------------------------------------------------------------------
+# PixelGen-XXL T2I SFT @ 512 + DMD2-style GAN branch.
+#
+# v8 change vs v7:
+# - Keep discriminator backbone trainable.
+# - Keep discriminator FM loss.
+# - Keep generator and discriminator LR fixed and equal, without decay.
+# - Use a fixed generator GAN loss weight of 0.003.
+# - D optimizes:
+# D loss = GAN discriminator loss + d_fm_loss_weight * D-backbone FM loss.
+# ---------------------------------------------------------------------------
+seed_everything: true
+tags:
+ exp: &exp sft_pix512_gan_v8
+torch_hub_dir: ~/.cache/torch/hub
+ckpt_path: /media/home/lx/PixelGen-fuxian/ck/PixelGen_XXL_T2I.ckpt
+huggingface_cache_dir: null
+trainer:
+ default_root_dir: ./universal_pix_t2i_workdirs
+ accelerator: auto
+ strategy: ddp_find_unused_parameters_true
+ devices: auto
+ num_nodes: 1
+ accumulate_grad_batches: 1
+ precision: bf16-mixed
+ logger:
+ class_path: lightning.pytorch.loggers.WandbLogger
+ init_args:
+ project: universal_flow_JiT_t2i
+ name: *exp
+ num_sanity_val_steps: 2
+ max_steps: 200000
+ val_check_interval: 10000
+ check_val_every_n_epoch: null
+ log_every_n_steps: 50
+ deterministic: null
+ inference_mode: true
+ use_distributed_sampler: false
+ callbacks:
+ - class_path: src.callbacks.model_checkpoint.CheckpointHook
+ init_args:
+ every_n_train_steps: 5000
+ save_top_k: -1
+ save_last: true
+ - class_path: src.callbacks.save_images.SaveImagesHook
+ init_args:
+ save_dir: val
+ plugins:
+ - src.plugins.bd_env.BDEnvironment
+model:
+ vae:
+ class_path: src.models.autoencoder.pixel.PixelAE
+ denoiser:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: &hidden_dim 1536
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: &txt_embed_dim 2048
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ conditioner:
+ class_path: src.models.conditioner.qwen3_text_encoder.Qwen3TextEncoder
+ init_args:
+ weight_path: Qwen/Qwen3-1.7B
+ embed_dim: *txt_embed_dim
+ max_length: 128
+ diffusion_trainer:
+ class_path: src.diffusion.flow_matching.training_repa_JiT_LPIPS_DINO_NoiseGating_DMD2_GAN.REPADMD2GANTrainer
+ init_args:
+ lognorm_t: true
+ encoder:
+ class_path: src.models.encoder.DINOv2
+ init_args:
+ weight_path: ~/.cache/torch/hub/checkpoints/dinov2_vitb14_pretrain.pth
+ align_layer: 8
+ proj_denoiser_dim: *hidden_dim
+ proj_hidden_dim: *hidden_dim
+ proj_encoder_dim: 768
+ P_mean: -0.8
+ P_std: 0.8
+ t_eps: 0.05
+ scheduler: &scheduler src.diffusion.flow_matching.scheduling.LinearScheduler
+ lpips_weight: 0.1
+ dino_weight: 0.01
+ percept_t_threshold: 0.3
+ percept_adaptive_step: 0
+ use_gan_loss: true
+ gan_loss_weight: 0.003
+ d_fm_loss_weight: 1.0
+ gan_loss_weight_final: 0.003
+ decay_start_step: 0
+ decay_end_step: 0
+ lr_final_multiplier: 1.0
+ decay_schedule: constant
+ diffusion_gan: true
+ # Add random flow-matching noise before D, DMD2-style. In this codebase
+ # t=1 is clean and t=0 is pure noise, so max 0.4 keeps D inputs in the
+ # same lower-noise region as gan_timestep_min=0.6.
+ diffusion_gan_max_timestep: 0.4
+ gan_start_step: 0
+ gan_timestep_min: 0.6
+ gan_timestep_max: 1.0
+ discriminator:
+ class_path: src.models.discriminator.DenoiserBackboneDiscriminator
+ init_args:
+ backbone:
+ class_path: src.models.transformer.JiT_T2I.JiT_T2I
+ init_args:
+ patch_size: 16
+ input_size: 512
+ in_channels: 3
+ hidden_size: *hidden_dim
+ num_blocks: 16
+ num_groups: 24
+ txt_embed_dim: *txt_embed_dim
+ txt_max_length: 128
+ bottleneck_dim: 256
+ num_text_blocks: 4
+ hidden_size: *hidden_dim
+ head_hidden_size: *hidden_dim
+ head_depth: 2
+ train_backbone: true
+ backbone_lr_scale: 1.0
+ d_optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.0
+ - 0.99
+ weight_decay: 0.00
+ g_grad_clip: 1.0
+ d_grad_clip: 1.0
+ diffusion_sampler:
+ class_path: src.diffusion.flow_matching.adam_sampling.AdamLMSamplerJiT
+ init_args:
+ num_steps: 25
+ guidance: 4.0
+ timeshift: 3.0
+ order: 2
+ scheduler: *scheduler
+ guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
+ ema_tracker:
+ class_path: src.callbacks.simple_ema.SimpleEMA
+ init_args:
+ decay: 0.9999
+ optimizer:
+ class_path: torch.optim.AdamW
+ init_args:
+ lr: 1.0e-5
+ betas:
+ - 0.9
+ - 0.999
+ weight_decay: 0.00
+data:
+ train_dataset:
+ class_path: src.data.dataset.image_txt.ImageText
+ init_args:
+ root: /media/home/lx/sharp/blip/extracted
+ resolution: 512
+ eval_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ limit: 50
+ seed_offset: 0
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ pred_dataset:
+ class_path: src.data.dataset.txt_dir_eval.TxtDirEvalDataset
+ init_args:
+ txt_dir: /media/home/lx/sharp/blip/extracted_new/sa_000201
+ limit: 50
+ seed_offset: 0
+ latent_shape:
+ - 3
+ - 512
+ - 512
+ train_batch_size: 8
+ train_num_workers: 8
+ pred_batch_size: 8
+ pred_num_workers: 1
diff --git a/docs/.ds_store b/docs/.ds_store
new file mode 100644
index 0000000000000000000000000000000000000000..4ea3de6ba04bb31e5fc03a91da40827f1e659952
Binary files /dev/null and b/docs/.ds_store differ
diff --git a/docs/index.html b/docs/index.html
new file mode 100644
index 0000000000000000000000000000000000000000..a7e9c6f691895cffff91130190872ebf6fb07ac0
--- /dev/null
+++ b/docs/index.html
@@ -0,0 +1,335 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ +
+
+
+
+
+
+
+
+
+ +
+ PixelGen: Pixel Diffusion Beats Latent Diffusion with Perceptual Loss +
+ + + + + + + +
+
+
+ 
+
+
+
+
+
+ Abstract
+ +
+
+
+
+ + Pixel diffusion generates images directly in pixel space in an end-to-end manner, avoiding the artifacts and bottlenecks introduced by VAEs in two-stage latent diffusion. However, it is challenging to optimize high-dimensional pixel manifolds that contain many perceptually irrelevant signals, leaving existing pixel diffusion methods lagging behind latent diffusion models. We propose PixelGen, a simple pixel diffusion framework with perceptual supervision. Instead of modeling the full image manifold, PixelGen introduces two complementary perceptual losses to guide diffusion model towards learning a more meaningful perceptual manifold. An LPIPS loss facilitates learning better local patterns, while a DINO-based perceptual loss strengthens global semantics. With perceptual supervision, PixelGen surpasses strong latent diffusion baselines. It achieves an FID of 5.11 on ImageNet-256 without classifier-free guidance using only 80 training epochs, and demonstrates favorable scaling performance on large-scale text-to-image generation with a GenEval score of 0.79. PixelGen requires no VAEs, no latent representations, and no auxiliary stages, providing a simpler yet more powerful generative paradigm. +
+
+
+
+
+
+
+
+ Motivation
+
+
+
+ + It is challenging to optimize high-dimensional pixel manifolds that contain many perceptually irrelevant signals, leaving existing pixel diffusion methods lagging behind latent diffusion models. We introduce PixelGen, a simple pixel diffusion framework with perceptual loss. Instead of modeling the full image manifold, PixelGen introduces two complementary perceptual losses to guide diffusion model towards learning a more meaningful perceptual manifold. An LPIPS loss facilitates learning better local patterns, while a DINO-based perceptual loss strengthens global semantics. +
+ +
+ 
+ Figure 2: This work shows that pixel diffusion with perceptual loss outperforms latent diffusion. (a) A traditional two-stage latent diffusion denoises in the latent space, which is influenced by the artifacts of the VAE. (b) PixelGen introduces perceptual loss to encourage the diffusion model to focus on the perceptual manifold, enabling the pixel diffusion to learn a meaningful manifold rather than the complex full image manifold. (c) PixelGen outperforms the latent diffusion models using only 80 training epochs on ImageNet without CFG. +
+ + Figure 2: This work shows that pixel diffusion with perceptual loss outperforms latent diffusion. (a) A traditional two-stage latent diffusion denoises in the latent space, which is influenced by the artifacts of the VAE. (b) PixelGen introduces perceptual loss to encourage the diffusion model to focus on the perceptual manifold, enabling the pixel diffusion to learn a meaningful manifold rather than the complex full image manifold. (c) PixelGen outperforms the latent diffusion models using only 80 training epochs on ImageNet without CFG. +
Illustration of Perceptual Manifold
+
+
+
+
+ 
+ Figure 3: Illustration of different manifolds within the pixel space. The image manifold is a large manifold containing both perceptually significant information and imperceptible signals. The perceptual manifold contains perceptually important signals, providing a better target for pixel space diffusion. P-DINO and LPIPS are the two complementary perceptual supervision utilized in PixelGen. +
+ + Figure 3: Illustration of different manifolds within the pixel space. The image manifold is a large manifold containing both perceptually significant information and imperceptible signals. The perceptual manifold contains perceptually important signals, providing a better target for pixel space diffusion. P-DINO and LPIPS are the two complementary perceptual supervision utilized in PixelGen. +
Implementation
+
+
+
+ + We propose PixelGen, a simple yet effective pixel diffusion framework with perceptual supervision. PixelGen directly operates in the pixel domain without relying on latent representations, VAEs, or auxiliary stages. Following the x-prediction paradigm, the diffusion model predicts clean images instead of noise or velocity. To retain the benefits of flow matching, the predicted image is converted into velocity, resulting in a flow-matching objective. + PixelGen focuses on the perceptual manifold rather than the full image manifold. To this end, we introduce two complementary perceptual losses. An LPIPS loss emphasizes local textures and fine-grained details, while a Perceptual DINO (P-DINO) loss aligns global semantics using patch-level features from a frozen DINOv2 encoder. +
+
+ 
+ Figure 4: Overview of PixelGen. The diffusion model directly predicts the image x instead of velocity or noise to simplify the prediction target. A flow-matching diffusion loss is retained to keep the advantages of flow matching via velocity conversion. Two complementary perceptual losses are introduced to encourage the diffusion model to focus on the perceptual manifold. +
+ + Figure 4: Overview of PixelGen. The diffusion model directly predicts the image x instead of velocity or noise to simplify the prediction target. A flow-matching diffusion loss is retained to keep the advantages of flow matching via velocity conversion. Two complementary perceptual losses are introduced to encourage the diffusion model to focus on the perceptual manifold. +
Empirically Analysis
+
+
+ + Perceptual supervision improves pixel diffusion by enhancing local details and global semantics. Starting from the JiT baseline, we progressively introduce the LPIPS loss and the P-DINO loss. As shown in Figure 5, the baseline model produces blurry images with weak structural consistency. After adding the LPIPS loss, local textures become sharper, and fine details are better preserved. This indicates that LPIPS effectively emphasizes perceptually important local patterns. When the P-DINO loss is further introduced, the generated images exhibit improved global structure and better semantics. These qualitative improvements are supported by quantitative results. The baseline model achieves an FID of 23.67 on ImageNet without classifier-free guidance. With LPIPS loss, the FID decreases to 10.00. After adding the P-DINO loss, it further drops to 7.46. This confirms that LPIPS and P-DINO provide complementary supervision. LPIPS focuses on local perceptual fidelity, while P-DINO enhances global semantics. Together, they guide the diffusion model toward a perceptually meaningful manifold. +
+
+ 
+ Figure 5: Effectiveness of perceptual supervision in PixelGen. LPIPS and P-DINO losses are progressively added to a baseline pixel diffusion model. The LPIPS loss improves local texture fidelity, while P-DINO further enhances global semantics. + +
+ + Figure 5: Effectiveness of perceptual supervision in PixelGen. LPIPS and P-DINO losses are progressively added to a baseline pixel diffusion model. The LPIPS loss improves local texture fidelity, while P-DINO further enhances global semantics. + +
+
+
+
+
+
+ Evaluations
+Quantitative Results
+
+ 
+
+
+
+
+
+
+
+
+
+ +
+
+ +
+
+ Qualitative Results
+
+ 
+ Figure 6: More Qualitative results of text-to-image generation at a 512x512 resolution. Our PixelGen supports multiple languages with the Qwen3 text encoder, such as Chinese and English.
+
+
+
+
+ Figure 7: Qualitative results of class-to-image generation at a 256x256 resolution.
+
+
+
+
+
+ Figure 6: More Qualitative results of text-to-image generation at a 512x512 resolution. Our PixelGen supports multiple languages with the Qwen3 text encoder, such as Chinese and English.
+
+ +
+
+ Figure 7: Qualitative results of class-to-image generation at a 256x256 resolution.
+
+ +
+
+BibTeX
+@misc{ma2025PixelGenfrequencyPixelGenupledpixeldiffusion,
+ title={PixelGen: Pixel Diffusion Beats Latent Diffusion with Perceptual Loss},
+ author={Zehong Ma and Ruihan Xu and Shiliang Zhang},
+ year={2025},
+ eprint={2602.02493},
+ archivePrefix={arXiv},
+ primaryClass={cs.CV},
+ url={https://arxiv.org/abs/2602.02493},
+ }