diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..380065f0fc7551add38dd5b38f3ebfacd860ffdf 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +unet/Z-Image/assets/Z-Image-Gallery.pdf filter=lfs diff=lfs merge=lfs -text +unet/Z-Image/assets/image_arena_all.jpg filter=lfs diff=lfs merge=lfs -text +unet/Z-Image/assets/decoupled-dmd.webp filter=lfs diff=lfs merge=lfs -text +unet/Z-Image/assets/architecture.webp filter=lfs diff=lfs merge=lfs -text +unet/Z-Image/assets/reasoning.png filter=lfs diff=lfs merge=lfs -text +unet/Z-Image/assets/DMDR.webp filter=lfs diff=lfs merge=lfs -text +unet/Z-Image/assets/showcase.jpg filter=lfs diff=lfs merge=lfs -text diff --git a/LLM/Florence-2-base/model.safetensors b/LLM/Florence-2-base/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..26fd400b0e459331dc9f96b37d6a3e629d59d66a --- /dev/null +++ b/LLM/Florence-2-base/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:03075d2d2d2bbd3e180b9ba0afae4aa8563226e2d32911656966e05b2f2ee060 +size 463221266 diff --git a/LLM/Florence-2-base/pytorch_model.bin b/LLM/Florence-2-base/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..435fb4d824c99715b72335e8ab80873df19fcfb6 --- /dev/null +++ b/LLM/Florence-2-base/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b480ac374593b0dcb18ffa63b23213734e04cd43eab0d620d23e39708d4a4a7e +size 464421827 diff --git a/SEEDVR2/ema_vae_fp16.safetensors b/SEEDVR2/ema_vae_fp16.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..289fac8e5236488d8b6d90884a3ff694a5c87578 --- /dev/null +++ b/SEEDVR2/ema_vae_fp16.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:20678548f420d98d26f11442d3528f8b8c94e57ee046ef93dbb7633da8612ca1 +size 501324814 diff --git a/assets/bpe_simple_vocab_16e6.txt.gz b/assets/bpe_simple_vocab_16e6.txt.gz new file mode 100644 index 0000000000000000000000000000000000000000..36a15856e00a06a9fbed8cdd34d2393fea4a3113 --- /dev/null +++ b/assets/bpe_simple_vocab_16e6.txt.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a +size 1356917 diff --git a/checkpoints/qwen_image_fp8_hq.safetensors b/checkpoints/qwen_image_fp8_hq.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5389c6f06575f00537b177c2ed4579159da9cdff --- /dev/null +++ b/checkpoints/qwen_image_fp8_hq.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5a61c42a58c181813fd94748df62ca1cdb53d0ec4b32c34af09375e5309126fa +size 89460 diff --git a/depthcrafter/stabilityai_stable-video-diffusion-img2vid-xt/vae/diffusion_pytorch_model.fp16.safetensors b/depthcrafter/stabilityai_stable-video-diffusion-img2vid-xt/vae/diffusion_pytorch_model.fp16.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..76c1b086f4ca2c1060a3ba7c8dea11d522a67584 --- /dev/null +++ b/depthcrafter/stabilityai_stable-video-diffusion-img2vid-xt/vae/diffusion_pytorch_model.fp16.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:af602cd0eb4ad6086ec94fbf1438dfb1be5ec9ac03fd0215640854e90d6463a3 +size 195531910 diff --git a/detection/vitpose_h_wholebody_model.onnx b/detection/vitpose_h_wholebody_model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..26819d54d1113adb18c975435f4c2104094802cd --- /dev/null +++ b/detection/vitpose_h_wholebody_model.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f21466cd6c93d0066782ad5923c14a4e6569133def212dc2895c73596c2e553b +size 420252 diff --git a/detection/yolov10m.onnx b/detection/yolov10m.onnx new file mode 100644 index 0000000000000000000000000000000000000000..b781a65ff19ca23d2db19972c662f71104a4ecc8 --- /dev/null +++ b/detection/yolov10m.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:89b526498a6d55f869a6ab52e3a2eb20ad45b3711c1f7de3dd9ca0b399dfd6d7 +size 61659339 diff --git a/mediapipe/selfie_multiclass_256x256.tflite b/mediapipe/selfie_multiclass_256x256.tflite new file mode 100644 index 0000000000000000000000000000000000000000..9ebdec318f4426502f8d825b8f0332c3e20e29b7 --- /dev/null +++ b/mediapipe/selfie_multiclass_256x256.tflite @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c6748b1253a99067ef71f7e26ca71096cd449baefa8f101900ea23016507e0e0 +size 16371837 diff --git a/sams/bpe_simple_vocab_16e6.txt.gz b/sams/bpe_simple_vocab_16e6.txt.gz new file mode 100644 index 0000000000000000000000000000000000000000..36a15856e00a06a9fbed8cdd34d2393fea4a3113 --- /dev/null +++ b/sams/bpe_simple_vocab_16e6.txt.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a +size 1356917 diff --git a/ultralytics/bbox/adetailerFootYolov8x_v20.pt b/ultralytics/bbox/adetailerFootYolov8x_v20.pt new file mode 100644 index 0000000000000000000000000000000000000000..6c9e3be927e050a70ca46608b947b5e9a57e07f5 --- /dev/null +++ b/ultralytics/bbox/adetailerFootYolov8x_v20.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9f39f32ab83b43002ca466605603b6c6dcff124ecfb23dab1c74c36ecb95cb4b +size 136712062 diff --git a/ultralytics/bbox/face_yolov8m.pt b/ultralytics/bbox/face_yolov8m.pt new file mode 100644 index 0000000000000000000000000000000000000000..7448998167649f4418ec3b3d8467903512b22ced --- /dev/null +++ b/ultralytics/bbox/face_yolov8m.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:717923c19b3f4bbf5250b728f1fa6b2cb72a33aed1d236ea9caf0e21ad943e5f +size 52026019 diff --git a/ultralytics/bbox/face_yolov8m[1].pt b/ultralytics/bbox/face_yolov8m[1].pt new file mode 100644 index 0000000000000000000000000000000000000000..7448998167649f4418ec3b3d8467903512b22ced --- /dev/null +++ b/ultralytics/bbox/face_yolov8m[1].pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:717923c19b3f4bbf5250b728f1fa6b2cb72a33aed1d236ea9caf0e21ad943e5f +size 52026019 diff --git a/unet/Z-Image/assets/DMDR.webp b/unet/Z-Image/assets/DMDR.webp new file mode 100644 index 0000000000000000000000000000000000000000..8613192a1e4f9026d165321eb98e2095e44e03e2 --- /dev/null +++ b/unet/Z-Image/assets/DMDR.webp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2e6f3053b98d097f2aa11d3892bd9307326db41b65336bea54dc5825a0e03077 +size 172966 diff --git a/unet/Z-Image/assets/Z-Image-Gallery.pdf b/unet/Z-Image/assets/Z-Image-Gallery.pdf new file mode 100644 index 0000000000000000000000000000000000000000..89f8a821673e81f57b858fdcda78e198bcc50e47 --- /dev/null +++ b/unet/Z-Image/assets/Z-Image-Gallery.pdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f9895b3246d2547bac74bbe0be975da500eaae93f2cad4248ad3281786b1ac6 +size 15767436 diff --git a/unet/Z-Image/assets/architecture.webp b/unet/Z-Image/assets/architecture.webp new file mode 100644 index 0000000000000000000000000000000000000000..1d34768cd18021cc9d64ebac954fa432ea0868c8 --- /dev/null +++ b/unet/Z-Image/assets/architecture.webp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:261af62ecc7e9749ae28e1d3a84e2f70a6c192d2017b7d8f020c7bff982ef59c +size 422392 diff --git a/unet/Z-Image/assets/decoupled-dmd.webp b/unet/Z-Image/assets/decoupled-dmd.webp new file mode 100644 index 0000000000000000000000000000000000000000..c1e5a7176d24a0bc0cc9faa0efe5e61791aafc79 --- /dev/null +++ b/unet/Z-Image/assets/decoupled-dmd.webp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4568ca559b997fc38f57dc1c3f5b1da3a3c144ae12419caa855ced972bf8c7aa +size 152062 diff --git a/unet/Z-Image/assets/image_arena_all.jpg b/unet/Z-Image/assets/image_arena_all.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e7a31f958ca4e096d5a0db57a6f7c360cd9ddab1 --- /dev/null +++ b/unet/Z-Image/assets/image_arena_all.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:899a87527d6fe44068bf1928dc7af60baefaca9b9566034e7ec0f5b15e5e3833 +size 1648993 diff --git a/unet/Z-Image/assets/reasoning.png b/unet/Z-Image/assets/reasoning.png new file mode 100644 index 0000000000000000000000000000000000000000..3a275a8df7ed7e120ea98a369a6e1a2911e658a6 --- /dev/null +++ b/unet/Z-Image/assets/reasoning.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:96c16b2c8d8dc67bb92ecc22d54b9955ab55136977f515bb76f4b2eb42eb3cdb +size 7696993 diff --git a/unet/Z-Image/assets/showcase.jpg b/unet/Z-Image/assets/showcase.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a55d1ce9217afe90f1791ed279e15c08437594c1 --- /dev/null +++ b/unet/Z-Image/assets/showcase.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f6ee74e066e00596e429f5a08140aebae1678e5935ce1e11ca6c1c6cd72432ee +size 6433750 diff --git a/unet/Z-Image/src/config/__init__.py b/unet/Z-Image/src/config/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70dd94bf004ede3f7f2953a9a82af98468ae276e --- /dev/null +++ b/unet/Z-Image/src/config/__init__.py @@ -0,0 +1,91 @@ +"""Z-Image Configuration.""" + +from .inference import ( + DEFAULT_CFG_TRUNCATION, + DEFAULT_GUIDANCE_SCALE, + DEFAULT_HEIGHT, + DEFAULT_INFERENCE_STEPS, + DEFAULT_MAX_SEQUENCE_LENGTH, + DEFAULT_WIDTH, +) +from .model import ( + ADALN_EMBED_DIM, + BASE_IMAGE_SEQ_LEN, + BASE_SHIFT, + BYTES_PER_GB, + DEFAULT_LOAD_DEVICE, + DEFAULT_LOAD_DTYPE_STR, + DEFAULT_SCHEDULER_NUM_TRAIN_TIMESTEPS, + DEFAULT_SCHEDULER_SHIFT, + DEFAULT_SCHEDULER_USE_DYNAMIC_SHIFTING, + DEFAULT_TRANSFORMER_CAP_FEAT_DIM, + DEFAULT_TRANSFORMER_DIM, + DEFAULT_TRANSFORMER_F_PATCH_SIZE, + DEFAULT_TRANSFORMER_IN_CHANNELS, + DEFAULT_TRANSFORMER_N_HEADS, + DEFAULT_TRANSFORMER_N_KV_HEADS, + DEFAULT_TRANSFORMER_N_LAYERS, + DEFAULT_TRANSFORMER_N_REFINER_LAYERS, + DEFAULT_TRANSFORMER_NORM_EPS, + DEFAULT_TRANSFORMER_PATCH_SIZE, + DEFAULT_TRANSFORMER_QK_NORM, + DEFAULT_TRANSFORMER_T_SCALE, + DEFAULT_VAE_IN_CHANNELS, + DEFAULT_VAE_LATENT_CHANNELS, + DEFAULT_VAE_NORM_NUM_GROUPS, + DEFAULT_VAE_OUT_CHANNELS, + DEFAULT_VAE_SCALE_FACTOR, + DEFAULT_VAE_SCALING_FACTOR, + FREQUENCY_EMBEDDING_SIZE, + MAX_IMAGE_SEQ_LEN, + MAX_PERIOD, + MAX_SHIFT, + ROPE_AXES_DIMS, + ROPE_AXES_LENS, + ROPE_THETA, + SEQ_MULTI_OF, +) + +__all__ = [ + "ADALN_EMBED_DIM", + "SEQ_MULTI_OF", + "ROPE_THETA", + "ROPE_AXES_DIMS", + "ROPE_AXES_LENS", + "FREQUENCY_EMBEDDING_SIZE", + "MAX_PERIOD", + "BASE_IMAGE_SEQ_LEN", + "MAX_IMAGE_SEQ_LEN", + "BASE_SHIFT", + "MAX_SHIFT", + "DEFAULT_VAE_SCALE_FACTOR", + "DEFAULT_VAE_IN_CHANNELS", + "DEFAULT_VAE_OUT_CHANNELS", + "DEFAULT_VAE_LATENT_CHANNELS", + "DEFAULT_VAE_NORM_NUM_GROUPS", + "DEFAULT_VAE_SCALING_FACTOR", + "DEFAULT_TRANSFORMER_PATCH_SIZE", + "DEFAULT_TRANSFORMER_F_PATCH_SIZE", + "DEFAULT_TRANSFORMER_IN_CHANNELS", + "DEFAULT_TRANSFORMER_DIM", + "DEFAULT_TRANSFORMER_N_LAYERS", + "DEFAULT_TRANSFORMER_N_REFINER_LAYERS", + "DEFAULT_TRANSFORMER_N_HEADS", + "DEFAULT_TRANSFORMER_N_KV_HEADS", + "DEFAULT_TRANSFORMER_NORM_EPS", + "DEFAULT_TRANSFORMER_QK_NORM", + "DEFAULT_TRANSFORMER_CAP_FEAT_DIM", + "DEFAULT_TRANSFORMER_T_SCALE", + "DEFAULT_SCHEDULER_NUM_TRAIN_TIMESTEPS", + "DEFAULT_SCHEDULER_SHIFT", + "DEFAULT_SCHEDULER_USE_DYNAMIC_SHIFTING", + "DEFAULT_LOAD_DEVICE", + "DEFAULT_LOAD_DTYPE_STR", + "BYTES_PER_GB", + "DEFAULT_HEIGHT", + "DEFAULT_WIDTH", + "DEFAULT_INFERENCE_STEPS", + "DEFAULT_GUIDANCE_SCALE", + "DEFAULT_CFG_TRUNCATION", + "DEFAULT_MAX_SEQUENCE_LENGTH", +] diff --git a/unet/Z-Image/src/config/inference.py b/unet/Z-Image/src/config/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..33adb47284ebd6e43e6445079bad0758ea7d54df --- /dev/null +++ b/unet/Z-Image/src/config/inference.py @@ -0,0 +1,8 @@ +"""Inference-specific configuration for Z-Image.""" + +DEFAULT_HEIGHT = 1024 +DEFAULT_WIDTH = 1024 +DEFAULT_INFERENCE_STEPS = 8 +DEFAULT_GUIDANCE_SCALE = 0.0 +DEFAULT_CFG_TRUNCATION = 1.0 +DEFAULT_MAX_SEQUENCE_LENGTH = 512 diff --git a/unet/Z-Image/src/config/manifests/z-image-turbo.txt b/unet/Z-Image/src/config/manifests/z-image-turbo.txt new file mode 100644 index 0000000000000000000000000000000000000000..ff3a8206c3af84f47069947ae13b6bc850985002 --- /dev/null +++ b/unet/Z-Image/src/config/manifests/z-image-turbo.txt @@ -0,0 +1,20 @@ +# Z-Image Model Manifest +# Format: +# Generated automatically - DO NOT edit manually + +5e3226ed72a9a4a080f2a4ca78b98ddc model_index.json +ca682fcc6c5a94cf726b7187e64b9411 scheduler/scheduler_config.json +1e97eb35d9d0b6aa60c58a8df8d7d99a text_encoder/config.json +30b85686b9a9b002e012494fadc027cb text_encoder/model-00001-of-00003.safetensors +e6a24ea164404a01ad2800dbae4e1a13 text_encoder/model-00002-of-00003.safetensors +09e190ed15ff14795b6277e023cfcb2d text_encoder/model-00003-of-00003.safetensors +589f5395156900f49d617aee8a8d8708 text_encoder/model.safetensors.index.json +6423133b9cc1a2077b57822c30c211aa tokenizer/tokenizer.json +b06e103ac555ec4b51266078b518c0f0 tokenizer/tokenizer_config.json +baed87136fe5f848e24b072f99856cc3 transformer/config.json +54889d0dd179b4fa2fd7bd0e487d856e transformer/diffusion_pytorch_model-00001-of-00003.safetensors +fe81e804658d345323512c63224b0604 transformer/diffusion_pytorch_model-00002-of-00003.safetensors +4e074e09129f98ad840414951f122feb transformer/diffusion_pytorch_model-00003-of-00003.safetensors +76d788eb0d42c59cc8f8ec007db639aa transformer/diffusion_pytorch_model.safetensors.index.json +ba9e2980c8630b4abccc643bc9f4a542 vae/config.json +6f83de55cb720c7fae051b14528577bf vae/diffusion_pytorch_model.safetensors diff --git a/unet/Z-Image/src/config/model.py b/unet/Z-Image/src/config/model.py new file mode 100644 index 0000000000000000000000000000000000000000..926033a358b9a54070ad78e7e27e40c0bf19abe2 --- /dev/null +++ b/unet/Z-Image/src/config/model.py @@ -0,0 +1,45 @@ +"""Model configuration constants for Z-Image.""" + +ADALN_EMBED_DIM = 256 +SEQ_MULTI_OF = 32 + +ROPE_THETA = 256.0 +ROPE_AXES_DIMS = [32, 48, 48] +ROPE_AXES_LENS = [1536, 512, 512] + +FREQUENCY_EMBEDDING_SIZE = 256 +MAX_PERIOD = 10000 + +BASE_IMAGE_SEQ_LEN = 256 +MAX_IMAGE_SEQ_LEN = 4096 +BASE_SHIFT = 0.5 +MAX_SHIFT = 1.15 + +DEFAULT_VAE_SCALE_FACTOR = 8 +DEFAULT_VAE_IN_CHANNELS = 3 +DEFAULT_VAE_OUT_CHANNELS = 3 +DEFAULT_VAE_LATENT_CHANNELS = 4 +DEFAULT_VAE_NORM_NUM_GROUPS = 32 +DEFAULT_VAE_SCALING_FACTOR = 0.18215 + +DEFAULT_TRANSFORMER_PATCH_SIZE = (2,) +DEFAULT_TRANSFORMER_F_PATCH_SIZE = (1,) +DEFAULT_TRANSFORMER_IN_CHANNELS = 16 +DEFAULT_TRANSFORMER_DIM = 3840 +DEFAULT_TRANSFORMER_N_LAYERS = 30 +DEFAULT_TRANSFORMER_N_REFINER_LAYERS = 2 +DEFAULT_TRANSFORMER_N_HEADS = 30 +DEFAULT_TRANSFORMER_N_KV_HEADS = 30 +DEFAULT_TRANSFORMER_NORM_EPS = 1e-5 +DEFAULT_TRANSFORMER_QK_NORM = True +DEFAULT_TRANSFORMER_CAP_FEAT_DIM = 2560 +DEFAULT_TRANSFORMER_T_SCALE = 1000.0 + +DEFAULT_SCHEDULER_NUM_TRAIN_TIMESTEPS = 1000 +DEFAULT_SCHEDULER_SHIFT = 3.0 +DEFAULT_SCHEDULER_USE_DYNAMIC_SHIFTING = False + +DEFAULT_LOAD_DEVICE = "cuda" +DEFAULT_LOAD_DTYPE_STR = "bfloat16" + +BYTES_PER_GB = 2**30 diff --git a/unet/Z-Image/src/tools/__init__.py b/unet/Z-Image/src/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..23504412375438174e900a28ef45768874d0ec41 --- /dev/null +++ b/unet/Z-Image/src/tools/__init__.py @@ -0,0 +1,9 @@ +"""Tools for Z-Image model management.""" + +from .generate_manifest import compute_md5, get_essential_files + +__all__ = [ + "compute_md5", + "get_essential_files", +] + diff --git a/unet/Z-Image/src/tools/generate_manifest.py b/unet/Z-Image/src/tools/generate_manifest.py new file mode 100644 index 0000000000000000000000000000000000000000..450571653583e73e3480f3d2ecd8909825fdedc1 --- /dev/null +++ b/unet/Z-Image/src/tools/generate_manifest.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""Generate manifest file with MD5 checksums for model weights. + +Usage: + python -m tools.generate_manifest ckpts/Z-Image-Turbo + python -m tools.generate_manifest ckpts/Z-Image-Turbo --no-checksums # Only list files +""" + +import argparse +import hashlib +from pathlib import Path +from typing import List + + +def compute_md5(file_path: Path, chunk_size: int = 8192) -> str: + """Compute MD5 hash of a file.""" + md5_hash = hashlib.md5() + with open(file_path, "rb") as f: + while chunk := f.read(chunk_size): + md5_hash.update(chunk) + return md5_hash.hexdigest() + + +def get_essential_files(model_dir: Path) -> List[Path]: + """Get list of essential model files.""" + essential_patterns = [ + "model_index.json", + "transformer/config.json", + "transformer/*.safetensors*", + "vae/config.json", + "vae/*.safetensors", + "text_encoder/config.json", + "text_encoder/*.safetensors*", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "scheduler/scheduler_config.json", + ] + + files = [] + for pattern in essential_patterns: + if "*" in pattern: + files.extend(model_dir.glob(pattern)) + else: + file_path = model_dir / pattern + if file_path.exists(): + files.append(file_path) + + return sorted(files) + + +def main(): + parser = argparse.ArgumentParser(description="Generate manifest file for model weights") + parser.add_argument("model_dir", type=str, help="Path to model directory") + parser.add_argument("--output", "-o", type=str, default=None, + help="Output manifest file path (default: auto-detect to config/manifests/)") + parser.add_argument("--no-checksums", action="store_true", + help="Only list files without computing checksums") + parser.add_argument("--verbose", "-v", action="store_true", + help="Print progress") + + args = parser.parse_args() + + model_dir = Path(args.model_dir) + if not model_dir.exists(): + print(f"Error: Model directory not found: {model_dir}") + return 1 + + # Determine output path + if args.output: + output_file = Path(args.output) + else: + # Auto-detect: save to config/manifests/{model-name}.txt + model_name = model_dir.name.lower() # e.g., "Z-Image-Turbo" -> "z-image-turbo" + script_dir = Path(__file__).parent + config_dir = script_dir.parent / "config" / "manifests" + config_dir.mkdir(parents=True, exist_ok=True) + output_file = config_dir / f"{model_name}.txt" + + # Get essential files + files = get_essential_files(model_dir) + + if not files: + print(f"Warning: No essential files found in {model_dir}") + return 1 + + print(f"Found {len(files)} essential files") + + # Generate manifest + with open(output_file, "w", encoding="utf-8") as f: + f.write("# Z-Image Model Manifest\n") + if args.no_checksums: + f.write("# Format: \n") + else: + f.write("# Format: \n") + f.write("# Generated automatically - DO NOT edit manually\n\n") + + for file_path in files: + rel_path = file_path.relative_to(model_dir) + + if args.no_checksums: + f.write(f"{rel_path}\n") + if args.verbose: + print(f" {rel_path}") + else: + if args.verbose: + print(f"Computing MD5 for {rel_path}...", end=" ", flush=True) + + try: + md5_hash = compute_md5(file_path) + f.write(f"{md5_hash} {rel_path}\n") + if args.verbose: + print(f"✓ {md5_hash}") + except Exception as e: + print(f"✗ Error: {e}") + continue + + print(f"\n✓ Manifest saved to: {output_file}") + print(f" Total files: {len(files)}") + if not args.no_checksums: + print(f" With MD5 checksums for integrity verification") + + return 0 + + +if __name__ == "__main__": + exit(main()) + diff --git a/unet/Z-Image/src/utils/__init__.py b/unet/Z-Image/src/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c65319ae2a1eba1e570221360fab56ea62df990e --- /dev/null +++ b/unet/Z-Image/src/utils/__init__.py @@ -0,0 +1,15 @@ +"""Utilities for Z-Image.""" + +from .attention import AttentionBackend, dispatch_attention, set_attention_backend +from .helpers import format_bytes, print_memory_stats, ensure_model_weights +from .loader import load_from_local_dir + +__all__ = [ + "load_from_local_dir", + "format_bytes", + "print_memory_stats", + "ensure_model_weights", + "AttentionBackend", + "set_attention_backend", + "dispatch_attention", +] diff --git a/unet/Z-Image/src/utils/attention.py b/unet/Z-Image/src/utils/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..768e5208cf1953a4b8621f865b98b72bafd34430 --- /dev/null +++ b/unet/Z-Image/src/utils/attention.py @@ -0,0 +1,516 @@ +"""Attention backend utilities for Z-Image.""" + +# Modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_dispatch.py +from enum import Enum +import functools +import inspect +from typing import Callable, Dict, List, Optional, Union + +import torch +import torch.nn.functional as F + +from .import_utils import is_flash_attn_3_available, is_flash_attn_available, is_torch_version + +_CAN_USE_FLASH_ATTN_2 = is_flash_attn_available() +_CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available() + +# MPS Flash Attention (Apple Silicon) +try: + import mps_flash_attn + _CAN_USE_MPS_FLASH = mps_flash_attn.is_available() +except ImportError: + _CAN_USE_MPS_FLASH = False + mps_flash_attn = None +_TORCH_VERSION_CHECK = is_torch_version(">=", "2.5.0") # have enable_gqa func call in SPDA + +if not _TORCH_VERSION_CHECK: + raise RuntimeError("PyTorch version must be >= 2.5.0 to use this backend.") +else: + print("PyTorch version is >= 2.5.0, check pass.") + +if _CAN_USE_FLASH_ATTN_2: + from flash_attn import flash_attn_func, flash_attn_varlen_func +else: + flash_attn_func = None + flash_attn_varlen_func = None + +if _CAN_USE_FLASH_ATTN_3: + from flash_attn_interface import ( + flash_attn_func as flash_attn_3_func, + flash_attn_varlen_func as flash_attn_3_varlen_func, + ) + + _flash_attn_3_sig = inspect.signature(flash_attn_3_func) + _FLASH_ATTN_3_SUPPORTS_RETURN_PROBS = "return_attn_probs" in _flash_attn_3_sig.parameters +else: + flash_attn_3_func = None + flash_attn_3_varlen_func = None + _FLASH_ATTN_3_SUPPORTS_RETURN_PROBS = False + + +class AttentionBackend(str, Enum): + """Supported attention backends.""" + + # Flash Attention + FLASH = "flash" + FLASH_VARLEN = "flash_varlen" + FLASH_3 = "_flash_3" + FLASH_VARLEN_3 = "_flash_varlen_3" + # MPS Flash Attention (Apple Silicon) + MPS_FLASH = "mps_flash" + # PyTorch Native Backends + NATIVE = "native" + NATIVE_FLASH = "_native_flash" + NATIVE_MATH = "_native_math" + + @classmethod + def print_available_backends(cls): + available_backends = [backend.value for backend in cls.__members__.values()] + print(f"Available attention backends list: {available_backends}") + + +# Registry for attention implementations +_ATTENTION_BACKENDS: Dict[str, Callable] = {} +_ATTENTION_CONSTRAINTS: Dict[str, List[Callable]] = {} + + +def register_backend(name: str, constraints: Optional[List[Callable]] = None): + def decorator(func): + _ATTENTION_BACKENDS[name] = func + _ATTENTION_CONSTRAINTS[name] = constraints or [] + return func + + return decorator + + +# --- Checks --- +def _check_device_cuda(query: torch.Tensor, **kwargs) -> None: + if query.device.type != "cuda": + raise ValueError("Query must be on a CUDA device.") + + +def _check_qkv_dtype_bf16_or_fp16(query: torch.Tensor, **kwargs) -> None: + if query.dtype not in (torch.bfloat16, torch.float16): + raise ValueError("Query must be either bfloat16 or float16.") + + +def _check_device_mps(query: torch.Tensor, **kwargs) -> None: + if query.device.type != "mps": + raise ValueError("Query must be on MPS device.") + + +def _process_mask(attn_mask: Optional[torch.Tensor], dtype: torch.dtype): + if attn_mask is None: + return None + + if attn_mask.ndim == 2: + attn_mask = attn_mask[:, None, None, :] + + # Convert bool mask to float additive mask + if attn_mask.dtype == torch.bool: + # NOTE: We skip checking for all-True mask (torch.all) to avoid graph breaks in torch.compile + new_mask = torch.zeros_like(attn_mask, dtype=dtype) + new_mask.masked_fill_(~attn_mask, float("-inf")) + return new_mask + + return attn_mask + + +def _normalize_attn_mask(attn_mask: torch.Tensor, batch_size: int, seq_len_k: int) -> torch.Tensor: + """Normalize an attention mask to shape [batch_size, seq_len_k] (bool).""" + if attn_mask.dtype != torch.bool: + # Try to convert float mask back to bool if possible, or assume it's float mask + # For varlen flash attn, we strictly need bool mask indicating valid tokens + if torch.is_floating_point(attn_mask): + return attn_mask > -1 # Assuming -inf is masked + # raise ValueError(f"Attention mask must be of type bool, got {attn_mask.dtype}.") + + if attn_mask.ndim == 1: + attn_mask = attn_mask.unsqueeze(0).expand(batch_size, seq_len_k) + elif attn_mask.ndim == 2: + if attn_mask.size(0) not in [1, batch_size]: + attn_mask = attn_mask.expand(batch_size, seq_len_k) + elif attn_mask.ndim == 3: + attn_mask = attn_mask.any(dim=1) + attn_mask = attn_mask.expand(batch_size, seq_len_k) + elif attn_mask.ndim == 4: + attn_mask = attn_mask.expand(batch_size, -1, -1, seq_len_k) + attn_mask = attn_mask.any(dim=(1, 2)) + + if attn_mask.shape != (batch_size, seq_len_k): + # Fallback reshape + return attn_mask.view(batch_size, seq_len_k) + + return attn_mask + + +@functools.lru_cache(maxsize=128) +def _prepare_for_flash_attn_varlen_without_mask( + batch_size: int, + seq_len_q: int, + seq_len_kv: int, + device: Optional[torch.device] = None, +): + # Optimized to avoid Inductor "pointless_cumsum_replacement" crash and remove graph breaks + seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device) + seqlens_k = torch.full((batch_size,), seq_len_kv, dtype=torch.int32, device=device) + + cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=device) * seq_len_q + cu_seqlens_k = torch.arange(batch_size + 1, dtype=torch.int32, device=device) * seq_len_kv + + return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (seq_len_q, seq_len_kv) + + +def _prepare_for_flash_attn_varlen_with_mask( + batch_size: int, + seq_len_q: int, + attn_mask: torch.Tensor, + device: Optional[torch.device] = None, +): + seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device) + seqlens_k = attn_mask.sum(dim=1, dtype=torch.int32) + # Use arange for Q to avoid Inductor crash + cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=device) * seq_len_q + + cu_seqlens_k = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0) + + max_seqlen_q = seq_len_q + max_seqlen_k = attn_mask.shape[1] # not max().item(), static shape to avoid graph break + + return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) + + +def _prepare_for_flash_attn_varlen( + batch_size: int, + seq_len_q: int, + seq_len_kv: int, + attn_mask: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, +) -> None: + if attn_mask is None: + return _prepare_for_flash_attn_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, device) + return _prepare_for_flash_attn_varlen_with_mask(batch_size, seq_len_q, attn_mask, device) + + +@register_backend(AttentionBackend.FLASH, constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16]) +def _flash_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, +) -> torch.Tensor: + if not _CAN_USE_FLASH_ATTN_2: + raise RuntimeError( + f"Flash Attention backend '{AttentionBackend.FLASH}' is not usable because of missing package." + ) + + out = flash_attn_func( + q=query, + k=key, + v=value, + dropout_p=dropout_p, + softmax_scale=scale, + causal=is_causal, + ) + return out + + +@register_backend(AttentionBackend.FLASH_VARLEN, constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16]) +def _flash_varlen_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, +) -> torch.Tensor: + if not _CAN_USE_FLASH_ATTN_2: + raise RuntimeError(f"Backend '{AttentionBackend.FLASH_VARLEN}' requires flash-attn.") + + batch_size, seq_len_q, _, _ = query.shape + _, seq_len_kv, _, _ = key.shape + + if attn_mask is not None: + attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) + + (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = _prepare_for_flash_attn_varlen( + batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device + ) + + query_packed = query.flatten(0, 1) + + if attn_mask is not None: + key_valid = [] + value_valid = [] + for b in range(batch_size): + valid_len = seqlens_k[b] + key_valid.append(key[b, :valid_len]) + value_valid.append(value[b, :valid_len]) + key_packed = torch.cat(key_valid, dim=0) + value_packed = torch.cat(value_valid, dim=0) + else: + key_packed = key.flatten(0, 1) + value_packed = value.flatten(0, 1) + + out = flash_attn_varlen_func( + q=query_packed, + k=key_packed, + v=value_packed, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + dropout_p=dropout_p, + softmax_scale=scale, + causal=is_causal, + ) + out = out.unflatten(0, (batch_size, -1)) + return out + + +@register_backend(AttentionBackend.FLASH_3, constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16]) +def _flash_attention_3( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, # Unused in simple FA3 func + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, +) -> torch.Tensor: + if not _CAN_USE_FLASH_ATTN_3: + raise RuntimeError(f"Backend '{AttentionBackend.FLASH_3}' requires Flash Attention 3 beta.") + + kwargs = { + "q": query, + "k": key, + "v": value, + "softmax_scale": scale, + "causal": is_causal, + } + + if _FLASH_ATTN_3_SUPPORTS_RETURN_PROBS: + kwargs["return_attn_probs"] = False + + out = flash_attn_3_func(**kwargs) + + if isinstance(out, tuple): + out = out[0] + + return out + + +@register_backend(AttentionBackend.FLASH_VARLEN_3, constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16]) +def _flash_varlen_attention_3( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, +) -> torch.Tensor: + if not _CAN_USE_FLASH_ATTN_3: + raise RuntimeError(f"Backend '{AttentionBackend.FLASH_VARLEN_3}' requires Flash Attention 3 beta.") + + batch_size, seq_len_q, _, _ = query.shape + _, seq_len_kv, _, _ = key.shape + + if attn_mask is not None: + attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) + + (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = _prepare_for_flash_attn_varlen( + batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device + ) + + query_packed = query.flatten(0, 1) + + if attn_mask is not None: + key_valid = [] + value_valid = [] + for b in range(batch_size): + valid_len = seqlens_k[b] + key_valid.append(key[b, :valid_len]) + value_valid.append(value[b, :valid_len]) + key_packed = torch.cat(key_valid, dim=0) + value_packed = torch.cat(value_valid, dim=0) + else: + key_packed = key.flatten(0, 1) + value_packed = value.flatten(0, 1) + + kwargs = { + "q": query_packed, + "k": key_packed, + "v": value_packed, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + "softmax_scale": scale, + "causal": is_causal, + } + + supports_return_probs = "return_attn_probs" in inspect.signature(flash_attn_3_varlen_func).parameters + + if supports_return_probs: + kwargs["return_attn_probs"] = False + + out = flash_attn_3_varlen_func(**kwargs) + + if isinstance(out, tuple): + out = out[0] + + out = out.unflatten(0, (batch_size, -1)) + return out + + +@register_backend(AttentionBackend.MPS_FLASH, constraints=[_check_device_mps, _check_qkv_dtype_bf16_or_fp16]) +def _mps_flash_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, +) -> torch.Tensor: + """MPS Flash Attention for Apple Silicon (M1/M2/M3/M4).""" + if not _CAN_USE_MPS_FLASH: + raise RuntimeError( + f"MPS Flash Attention backend '{AttentionBackend.MPS_FLASH}' requires mps-flash-attn package. " + "Install with: pip install mps-flash-attn" + ) + + # Convert from (B, S, H, D) to (B, H, S, D) for mps-flash-attn + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + # Convert mask to MFA format (bool, True = masked) + mfa_mask = None + if attn_mask is not None: + mfa_mask = mps_flash_attn.convert_mask(_process_mask(attn_mask, query.dtype)) + + out = mps_flash_attn.flash_attention( + query, key, value, + is_causal=is_causal, + scale=scale, + attn_mask=mfa_mask, + ) + + # Convert back to (B, S, H, D) + return out.transpose(1, 2).contiguous() + + +def _native_attention_wrapper( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + backend_kernel=None, +) -> torch.Tensor: + + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + attn_mask = _process_mask(attn_mask, query.dtype) + + if backend_kernel is not None: + with torch.nn.attention.sdpa_kernel(backend_kernel): + out = F.scaled_dot_product_attention( + query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale + ) + else: + out = F.scaled_dot_product_attention( + query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale + ) + + return out.transpose(1, 2).contiguous() + + +@register_backend(AttentionBackend.NATIVE_FLASH) +def _native_flash_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, +) -> torch.Tensor: + return _native_attention_wrapper( + query, + key, + value, + attn_mask=None, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + backend_kernel=torch.nn.attention.SDPBackend.FLASH_ATTENTION, + ) + + +@register_backend(AttentionBackend.NATIVE_MATH) +def _math_attention(*args, **kwargs): + return _native_attention_wrapper(*args, **kwargs, backend_kernel=torch.nn.attention.SDPBackend.MATH) + + +@register_backend(AttentionBackend.NATIVE) +def _native_attention(*args, **kwargs): + return _native_attention_wrapper(*args, **kwargs, backend_kernel=None) + + +def dispatch_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[float] = None, + backend: Union[str, AttentionBackend, None] = None, +) -> torch.Tensor: + + if isinstance(backend, AttentionBackend): + backend = backend.value + elif backend is None: + backend = AttentionBackend.NATIVE + else: + backend = str(backend) + + # Explicit dispatch to avoid dynamo guard issues on global dict + if backend == AttentionBackend.FLASH: + return _flash_attention(query, key, value, attn_mask, dropout_p, is_causal, scale) + elif backend == AttentionBackend.FLASH_VARLEN: + return _flash_varlen_attention(query, key, value, attn_mask, dropout_p, is_causal, scale) + elif backend == AttentionBackend.FLASH_3: + return _flash_attention_3(query, key, value, attn_mask, dropout_p, is_causal, scale) + elif backend == AttentionBackend.FLASH_VARLEN_3: + return _flash_varlen_attention_3(query, key, value, attn_mask, dropout_p, is_causal, scale) + elif backend == AttentionBackend.MPS_FLASH: + return _mps_flash_attention(query, key, value, attn_mask, dropout_p, is_causal, scale) + elif backend == AttentionBackend.NATIVE_FLASH: + return _native_flash_attention(query, key, value, attn_mask, dropout_p, is_causal, scale) + elif backend == AttentionBackend.NATIVE_MATH: + return _math_attention(query, key, value, attn_mask, dropout_p, is_causal, scale) + else: + return _native_attention(query, key, value, attn_mask, dropout_p, is_causal, scale) + + +def set_attention_backend(backend: Union[str, AttentionBackend, None]): + try: + from zimage.transformer import ZImageAttention + + if backend is not None: + backend = str(backend) + ZImageAttention._attention_backend = backend + except ImportError: + pass diff --git a/unet/Z-Image/src/utils/helpers.py b/unet/Z-Image/src/utils/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..15b78f579acc07537a6d65249dc2339adf02d4ad --- /dev/null +++ b/unet/Z-Image/src/utils/helpers.py @@ -0,0 +1,260 @@ +"""Helper utilities for Z-Image.""" + +import hashlib +import json +from pathlib import Path +from typing import Optional, List, Tuple, Dict + +from loguru import logger +import torch + +from config import BYTES_PER_GB + + +def format_bytes(size: float) -> str: + """ + Format bytes to GB string. + + Args: + size: Size in bytes + + Returns: + Formatted string in GB + """ + n = size / BYTES_PER_GB + return f"{n:.2f} GB" + + +def print_memory_stats(stage: str) -> None: + """ + Print CUDA memory statistics. + + Args: + stage: Description of current stage + """ + if not torch.cuda.is_available(): + logger.warning("CUDA not available, skipping memory stats") + return + + torch.cuda.synchronize() + allocated = torch.cuda.max_memory_allocated() + reserved = torch.cuda.max_memory_reserved() + current_allocated = torch.cuda.memory_allocated() + current_reserved = torch.cuda.memory_reserved() + + logger.info(f"[{stage}] Memory Stats:") + logger.info(f" Current Allocated: {format_bytes(current_allocated)}") + logger.info(f" Current Reserved: {format_bytes(current_reserved)}") + logger.info(f" Peak Allocated: {format_bytes(allocated)}") + logger.info(f" Peak Reserved: {format_bytes(reserved)}") + + +def compute_file_md5(file_path: Path, chunk_size: int = 8192) -> str: + """Compute MD5 hash of a file.""" + md5_hash = hashlib.md5() + with open(file_path, "rb") as f: + while chunk := f.read(chunk_size): + md5_hash.update(chunk) + return md5_hash.hexdigest() + + +def load_manifest(manifest_file: Path) -> Dict[str, Optional[str]]: + """Load manifest file. Returns dict mapping file paths to MD5 hashes (or None).""" + manifest = {} + if not manifest_file.exists(): + return manifest + + with open(manifest_file, "r", encoding="utf-8") as f: + for line_num, line in enumerate(f, 1): + line = line.strip() + # Skip empty lines and comments + if not line or line.startswith("#"): + continue + + parts = line.split() + + if len(parts) == 1: + # Only file path, no checksum + file_path = parts[0] + manifest[file_path] = None + elif len(parts) == 2: + # File path with checksum + if len(parts[0]) == 32 and all(c in '0123456789abcdef' for c in parts[0].lower()): + md5_hash, file_path = parts + else: + file_path, md5_hash = parts + manifest[file_path] = md5_hash + else: + logger.warning(f"Invalid manifest format at line {line_num}: {line}") + continue + + return manifest + + +def verify_file_integrity( + base_dir: Path, + manifest: Dict[str, Optional[str]], + verify_checksums: bool = True +) -> Tuple[bool, List[str], List[str]]: + """ + Verify file integrity using a manifest. + + Args: + base_dir: Base directory for relative file paths + manifest: Dictionary of relative paths to MD5 hashes (None if no hash provided) + verify_checksums: If True, verify MD5 checksums when available; if False, only check existence + + Returns: + Tuple of (all_valid: bool, missing_files: List[str], corrupted_files: List[str]) + """ + missing = [] + corrupted = [] + + for rel_path, expected_md5 in manifest.items(): + file_path = base_dir / rel_path + + if not file_path.exists(): + missing.append(rel_path) + continue + + # Only verify checksum if requested AND hash is available + if verify_checksums and expected_md5 is not None: + try: + actual_md5 = compute_file_md5(file_path) + if actual_md5 != expected_md5: + corrupted.append(rel_path) + logger.debug(f"Checksum mismatch for {rel_path}: expected {expected_md5}, got {actual_md5}") + except Exception as e: + logger.error(f"Failed to compute checksum for {rel_path}: {e}") + corrupted.append(rel_path) + + all_valid = len(missing) == 0 and len(corrupted) == 0 + return all_valid, missing, corrupted + + +def ensure_model_weights( + model_path: str, + repo_id: str = "Tongyi-MAI/Z-Image-Turbo", + verify: bool = False, + manifest_name: Optional[str] = None +) -> Path: + """ + Ensure model weights exist and optionally verify integrity. + + Args: + model_path: Path to model directory + repo_id: HuggingFace repo ID for download + verify: If True, verify MD5 checksums; if False, only check existence + manifest_name: Manifest file name in src/config/manifests/ (auto-detect if None) + + Returns: + Path to validated model directory + """ + from huggingface_hub import snapshot_download + + target_dir = Path(model_path) + + # Determine manifest path + if manifest_name: + # Explicitly specified manifest from config/manifests/ + manifest_path = Path(__file__).parent.parent / "config" / "manifests" / manifest_name + else: + # Auto-detect + model_name = target_dir.name.lower() # e.g., "Z-Image-Turbo" -> "z-image-turbo" + config_manifest = Path(__file__).parent.parent / "config" / "manifests" / f"{model_name}.txt" + + if config_manifest.exists(): + manifest_path = config_manifest + else: + # Fallback + manifest_path = target_dir / "manifest.txt" + + manifest = load_manifest(manifest_path) + + if not manifest: + logger.warning(f"Manifest file not found: {manifest_path}") + logger.warning("Skipping file verification (assuming model exists)") + if target_dir.exists(): + logger.info(f"✓ Model directory exists: {target_dir}") + return target_dir + else: + logger.warning(f"Model directory not found: {target_dir}") + missing_files = ["entire model directory"] + corrupted_files = [] + else: + # Count files with checksums + files_with_checksums = sum(1 for v in manifest.values() if v is not None) + + if verify and files_with_checksums == 0: + logger.info(f"Verify requested but no checksums in manifest, only checking existence") + elif verify and files_with_checksums > 0: + logger.info(f"Verifying {files_with_checksums} file(s) with MD5 checksums...") + + # Verify files + all_valid, missing_files, corrupted_files = verify_file_integrity( + target_dir, manifest, verify_checksums=verify + ) + + if all_valid: + if verify and files_with_checksums > 0: + logger.success(f"✓ All files verified with MD5 checksums in {target_dir}") + else: + logger.info(f"✓ All {len(manifest)} required files exist in {target_dir}") + return target_dir + + # Report missing and corrupted files + if missing_files: + logger.warning(f"Missing {len(missing_files)} file(s):") + for f in missing_files[:10]: + logger.warning(f" - {f}") + if len(missing_files) > 10: + logger.warning(f" ... and {len(missing_files) - 10} more") + + if corrupted_files: + logger.error(f"Corrupted {len(corrupted_files)} file(s) (checksum mismatch):") + for f in corrupted_files[:10]: + logger.error(f" - {f}") + if len(corrupted_files) > 10: + logger.error(f" ... and {len(corrupted_files) - 10} more") + + # Download model weights + logger.info(f"\nAttempting to download from {repo_id}...") + try: + target_dir.mkdir(parents=True, exist_ok=True) + snapshot_download( + repo_id=repo_id, + local_dir=str(target_dir), + local_dir_use_symlinks=False, + resume_download=True, + ) + logger.success("✓ Download completed") + except Exception as e: + logger.error(f"✗ Download failed: {e}") + logger.info( + f"\nIf you are offline, please manually download from:\n" + f" https://huggingface.co/{repo_id}\n" + f"and place in: {target_dir.absolute()}" + ) + raise RuntimeError(f"Failed to download model weights: {e}") + + # Verify after download + if manifest: + all_valid, missing_after, corrupted_after = verify_file_integrity( + target_dir, manifest, verify_checksums=verify + ) + + if not all_valid: + error_msg = [] + if missing_after: + error_msg.append(f"Still missing {len(missing_after)} file(s)") + if corrupted_after: + error_msg.append(f"Still corrupted {len(corrupted_after)} file(s)") + + raise FileNotFoundError( + f"After download: {', '.join(error_msg)}\n" + f"Please verify the download or manually place files in:\n" + f" {target_dir.absolute()}" + ) + + logger.success("✓ All model weights validated successfully") + return target_dir diff --git a/unet/Z-Image/src/utils/import_utils.py b/unet/Z-Image/src/utils/import_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6d64a30cca6a09a6f7f56f35f38f77775d6a6ee0 --- /dev/null +++ b/unet/Z-Image/src/utils/import_utils.py @@ -0,0 +1,31 @@ +import importlib.util + +import torch + + +def is_flash_attn_available(): + return importlib.util.find_spec("flash_attn") is not None + + +def is_flash_attn_3_available(): + return importlib.util.find_spec("flash_attn_interface") is not None + + +def is_torch_version(operator: str, version: str): + from packaging import version as pversion + + torch_version = pversion.parse(torch.__version__) + target_version = pversion.parse(version) + + # print(f"torch_version: {torch_version}, target: torch{operator}{target_version}") + if operator == ">": + return torch_version > target_version + elif operator == ">=": + return torch_version >= target_version + elif operator == "==": + return torch_version == target_version + elif operator == "<=": + return torch_version <= target_version + elif operator == "<": + return torch_version < target_version + return False diff --git a/unet/Z-Image/src/utils/loader.py b/unet/Z-Image/src/utils/loader.py new file mode 100644 index 0000000000000000000000000000000000000000..22f44eb0a8b49bd17fdd10ca4c7a96e835cfe68b --- /dev/null +++ b/unet/Z-Image/src/utils/loader.py @@ -0,0 +1,224 @@ +"""Model loading utilities for Z-Image components.""" + +import json +import os +from pathlib import Path +import sys +from typing import Optional, Union + +from loguru import logger +from safetensors.torch import load_file +import torch +from transformers import AutoModel, AutoTokenizer + +from config import ( + DEFAULT_SCHEDULER_NUM_TRAIN_TIMESTEPS, + DEFAULT_SCHEDULER_SHIFT, + DEFAULT_SCHEDULER_USE_DYNAMIC_SHIFTING, + DEFAULT_TRANSFORMER_CAP_FEAT_DIM, + DEFAULT_TRANSFORMER_DIM, + DEFAULT_TRANSFORMER_F_PATCH_SIZE, + DEFAULT_TRANSFORMER_IN_CHANNELS, + DEFAULT_TRANSFORMER_N_HEADS, + DEFAULT_TRANSFORMER_N_KV_HEADS, + DEFAULT_TRANSFORMER_N_LAYERS, + DEFAULT_TRANSFORMER_N_REFINER_LAYERS, + DEFAULT_TRANSFORMER_NORM_EPS, + DEFAULT_TRANSFORMER_PATCH_SIZE, + DEFAULT_TRANSFORMER_QK_NORM, + DEFAULT_TRANSFORMER_T_SCALE, + DEFAULT_VAE_IN_CHANNELS, + DEFAULT_VAE_LATENT_CHANNELS, + DEFAULT_VAE_NORM_NUM_GROUPS, + DEFAULT_VAE_OUT_CHANNELS, + DEFAULT_VAE_SCALING_FACTOR, + ROPE_AXES_DIMS, + ROPE_AXES_LENS, + ROPE_THETA, +) +from zimage.autoencoder import AutoencoderKL as LocalAutoencoderKL +from zimage.scheduler import FlowMatchEulerDiscreteScheduler + +DIFFUSERS_AVAILABLE = False + + +def load_config(config_path: str) -> dict: + with open(config_path, "r") as f: + return json.load(f) + + +def load_sharded_safetensors(weight_dir: Path, device: str = "cuda", dtype: Optional[torch.dtype] = None) -> dict: + """Load sharded safetensors from a directory.""" + weight_dir = Path(weight_dir) + index_files = list(weight_dir.glob("*.safetensors.index.json")) + + state_dict = {} + if index_files: + # Load sharded weights + with open(index_files[0], "r") as f: + index = json.load(f) + weight_map = index.get("weight_map", {}) + shard_files = set(weight_map.values()) + for shard_file in shard_files: + shard_path = weight_dir / shard_file + shard_state = load_file(str(shard_path), device=str(device)) + state_dict.update(shard_state) + else: + # Load single safetensors file + safetensors_files = list(weight_dir.glob("*.safetensors")) + if not safetensors_files: + raise FileNotFoundError(f"No safetensors files found in {weight_dir}") + state_dict = load_file(str(safetensors_files[0]), device=str(device)) + + # Cast to target dtype if specified + if dtype is not None: + state_dict = {k: v.to(dtype) if v.dtype != dtype else v for k, v in state_dict.items()} + + return state_dict + + +def load_from_local_dir( + model_dir: Union[str, Path], + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + verbose: bool = False, + compile: bool = False, +) -> dict: + """ + Load all Z-Image components from local directory. + + Args: + model_dir: Path to model directory + device: Device to load models on + dtype: Data type for model weights + verbose: Whether to display loading logs + compile: Whether to compile transformer and vae with torch.compile + + Returns: + Dictionary containing transformer, vae, text_encoder, tokenizer, and scheduler + """ + model_dir = Path(model_dir) + + sys.path.insert(0, str(model_dir.parent.parent / "Z-Image" / "src")) + from zimage.transformer import ZImageTransformer2DModel + + if verbose: + logger.info(f"Loading Z-Image from: {model_dir}") + + # DiT + if verbose: + logger.info("Loading DiT...") + transformer_dir = model_dir / "transformer" + config = load_config(str(transformer_dir / "config.json")) + + with torch.device("meta"): + transformer = ZImageTransformer2DModel( + all_patch_size=tuple(config.get("all_patch_size", DEFAULT_TRANSFORMER_PATCH_SIZE)), + all_f_patch_size=tuple(config.get("all_f_patch_size", DEFAULT_TRANSFORMER_F_PATCH_SIZE)), + in_channels=config.get("in_channels", DEFAULT_TRANSFORMER_IN_CHANNELS), + dim=config.get("dim", DEFAULT_TRANSFORMER_DIM), + n_layers=config.get("n_layers", DEFAULT_TRANSFORMER_N_LAYERS), + n_refiner_layers=config.get("n_refiner_layers", DEFAULT_TRANSFORMER_N_REFINER_LAYERS), + n_heads=config.get("n_heads", DEFAULT_TRANSFORMER_N_HEADS), + n_kv_heads=config.get("n_kv_heads", DEFAULT_TRANSFORMER_N_KV_HEADS), + norm_eps=config.get("norm_eps", DEFAULT_TRANSFORMER_NORM_EPS), + qk_norm=config.get("qk_norm", DEFAULT_TRANSFORMER_QK_NORM), + cap_feat_dim=config.get("cap_feat_dim", DEFAULT_TRANSFORMER_CAP_FEAT_DIM), + rope_theta=config.get("rope_theta", ROPE_THETA), + t_scale=config.get("t_scale", DEFAULT_TRANSFORMER_T_SCALE), + axes_dims=config.get("axes_dims", ROPE_AXES_DIMS), + axes_lens=config.get("axes_lens", ROPE_AXES_LENS), + ).to(dtype) + + # DiT (weights to CPU then move to GPU to optimize memory) + state_dict = load_sharded_safetensors(transformer_dir, device="cpu", dtype=dtype) + transformer.load_state_dict(state_dict, strict=False, assign=True) + del state_dict + + if verbose: + logger.info("Moving DiT to GPU...") + transformer = transformer.to(device) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + transformer.eval() + + # VAE + if verbose: + logger.info("Loading VAE...") + vae_dir = model_dir / "vae" + vae_config = load_config(str(vae_dir / "config.json")) + + vae = LocalAutoencoderKL( + in_channels=vae_config.get("in_channels", DEFAULT_VAE_IN_CHANNELS), + out_channels=vae_config.get("out_channels", DEFAULT_VAE_OUT_CHANNELS), + down_block_types=tuple(vae_config.get("down_block_types", ("DownEncoderBlock2D",))), + up_block_types=tuple(vae_config.get("up_block_types", ("UpDecoderBlock2D",))), + block_out_channels=tuple(vae_config.get("block_out_channels", (64,))), + layers_per_block=vae_config.get("layers_per_block", 1), + latent_channels=vae_config.get("latent_channels", DEFAULT_VAE_LATENT_CHANNELS), + norm_num_groups=vae_config.get("norm_num_groups", DEFAULT_VAE_NORM_NUM_GROUPS), + scaling_factor=vae_config.get("scaling_factor", DEFAULT_VAE_SCALING_FACTOR), + shift_factor=vae_config.get("shift_factor", None), + use_quant_conv=vae_config.get("use_quant_conv", True), + use_post_quant_conv=vae_config.get("use_post_quant_conv", True), + mid_block_add_attention=vae_config.get("mid_block_add_attention", True), + ) + + # VAE (fp32 for better precision) + vae_state_dict = load_sharded_safetensors(vae_dir, device="cpu") + vae.load_state_dict(vae_state_dict, strict=False) + del vae_state_dict + vae.to(device=device, dtype=torch.float32) + vae.eval() + torch.cuda.empty_cache() + + # Text Encoder + if verbose: + logger.info("Loading Text Encoder...") + text_encoder_dir = model_dir / "text_encoder" + text_encoder = AutoModel.from_pretrained( + str(text_encoder_dir), + # torch_dtype=dtype, # some version use this + dtype=dtype, + trust_remote_code=True, + ) + text_encoder.to(device) + text_encoder.eval() + + # Tokenizer + os.environ["TOKENIZERS_PARALLELISM"] = "false" + if verbose: + logger.info("Loading Tokenizer...") + tokenizer_dir = model_dir / "tokenizer" + tokenizer = AutoTokenizer.from_pretrained( + str(tokenizer_dir) if tokenizer_dir.exists() else str(text_encoder_dir), + trust_remote_code=True, + ) + + # Scheduler + if verbose: + logger.info("Loading Scheduler...") + scheduler_dir = model_dir / "scheduler" + scheduler_config = load_config(str(scheduler_dir / "scheduler_config.json")) + scheduler = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=scheduler_config.get("num_train_timesteps", DEFAULT_SCHEDULER_NUM_TRAIN_TIMESTEPS), + shift=scheduler_config.get("shift", DEFAULT_SCHEDULER_SHIFT), + use_dynamic_shifting=scheduler_config.get("use_dynamic_shifting", DEFAULT_SCHEDULER_USE_DYNAMIC_SHIFTING), + ) + + if compile: + if verbose: + logger.info("Compiling DiT and VAE...") + transformer = torch.compile(transformer) + vae = torch.compile(vae) + + if verbose: + logger.success("All components loaded successfully") + + return { + "transformer": transformer, + "vae": vae, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "scheduler": scheduler, + } diff --git a/unet/Z-Image/src/zimage/__init__.py b/unet/Z-Image/src/zimage/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b18cfb20e8803ae60fdaf886ffe15676d0129c69 --- /dev/null +++ b/unet/Z-Image/src/zimage/__init__.py @@ -0,0 +1,9 @@ +"""Z-Image PyTorch Native Implementation.""" + +from .pipeline import generate +from .transformer import ZImageTransformer2DModel + +__all__ = [ + "ZImageTransformer2DModel", + "generate", +] diff --git a/unet/Z-Image/src/zimage/autoencoder.py b/unet/Z-Image/src/zimage/autoencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..52062f4fa76e246875a522350c1f1f9bef917deb --- /dev/null +++ b/unet/Z-Image/src/zimage/autoencoder.py @@ -0,0 +1,369 @@ +"""AutoencoderKL implementation compatible with diffusers weights.""" + +# Modified from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/autoencoder.py +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn + + +@dataclass +class AutoencoderKLOutput: + sample: torch.Tensor + + +class AutoencoderConfig: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + def get(self, key, default=None): + return self.__dict__.get(key, default) + + def __getattr__(self, name): + return self.__dict__.get(name) + + +def swish(x): + return x * torch.sigmoid(x) + + +class ResnetBlock2D(nn.Module): + def __init__(self, in_channels, out_channels=None, dropout=0.0, temb_channels=512, groups=32, eps=1e-6): + super().__init__() + out_channels = out_channels or in_channels + self.in_channels = in_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True) + self.dropout = nn.Dropout(dropout) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + + self.nonlinearity = swish + + if self.in_channels != self.out_channels: + self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + else: + self.conv_shortcut = None + + def forward(self, input_tensor, temb=None): + hidden_states = input_tensor + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = (input_tensor + hidden_states) / 1.0 + return output_tensor + + +class Attention(nn.Module): + def __init__(self, in_channels, heads=1, dim_head=None, groups=32, eps=1e-6): + super().__init__() + self.heads = heads + self.in_channels = in_channels + self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + + self.to_q = nn.Linear(in_channels, in_channels) + self.to_k = nn.Linear(in_channels, in_channels) + self.to_v = nn.Linear(in_channels, in_channels) + self.to_out = nn.ModuleList([nn.Linear(in_channels, in_channels)]) + + def forward(self, hidden_states): + b, c, h, w = hidden_states.shape + residual = hidden_states + hidden_states = self.group_norm(hidden_states) + hidden_states = hidden_states.view(b, c, -1).transpose(1, 2) # (B, H*W, C) + + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + import torch.nn.functional as F + + hidden_states = F.scaled_dot_product_attention(query, key, value) + + hidden_states = self.to_out[0](hidden_states) + hidden_states = hidden_states.transpose(1, 2).view(b, c, h, w) + + return residual + hidden_states + + +class Downsample2D(nn.Module): + def __init__(self, channels, with_conv=True, out_channels=None, padding=1): + super().__init__() + out_channels = out_channels or channels + self.with_conv = with_conv + if with_conv: + self.conv = nn.Conv2d(channels, out_channels, kernel_size=3, stride=2, padding=padding) + + def forward(self, hidden_states): + if self.with_conv: + return self.conv(hidden_states) + else: + return torch.nn.functional.avg_pool2d(hidden_states, kernel_size=2, stride=2) + + +class Upsample2D(nn.Module): + def __init__(self, channels, with_conv=True, out_channels=None): + super().__init__() + out_channels = out_channels or channels + self.with_conv = with_conv + if with_conv: + self.conv = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, hidden_states): + hidden_states = torch.nn.functional.interpolate(hidden_states, scale_factor=2.0, mode="nearest") + if self.with_conv: + hidden_states = self.conv(hidden_states) + return hidden_states + + +class DownEncoderBlock2D(nn.Module): + def __init__(self, in_channels, out_channels, num_layers=1, resnet_eps=1e-6, resnet_groups=32, add_downsample=True): + super().__init__() + resnets = [] + for i in range(num_layers): + in_c = in_channels if i == 0 else out_channels + resnets.append(ResnetBlock2D(in_c, out_channels, eps=resnet_eps, groups=resnet_groups)) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [Downsample2D(out_channels, with_conv=True, out_channels=out_channels, padding=0)] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states): + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + pad = (0, 1, 0, 1) + hidden_states = torch.nn.functional.pad(hidden_states, pad, mode="constant", value=0) + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class UpDecoderBlock2D(nn.Module): + def __init__(self, in_channels, out_channels, num_layers=1, resnet_eps=1e-6, resnet_groups=32, add_upsample=True): + super().__init__() + resnets = [] + for i in range(num_layers): + in_c = in_channels if i == 0 else out_channels + resnets.append(ResnetBlock2D(in_c, out_channels, eps=resnet_eps, groups=resnet_groups)) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, with_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + def forward(self, hidden_states): + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class UNetMidBlock2D(nn.Module): + def __init__(self, in_channels, resnet_eps=1e-6, resnet_groups=32, attention_head_dim=None): + super().__init__() + self.resnets = nn.ModuleList( + [ + ResnetBlock2D(in_channels, in_channels, eps=resnet_eps, groups=resnet_groups), + ResnetBlock2D(in_channels, in_channels, eps=resnet_eps, groups=resnet_groups), + ] + ) + self.attentions = nn.ModuleList([Attention(in_channels, heads=1, groups=resnet_groups, eps=resnet_eps)]) + + def forward(self, hidden_states): + hidden_states = self.resnets[0](hidden_states) + for attn in self.attentions: + hidden_states = attn(hidden_states) + hidden_states = self.resnets[1](hidden_states) + return hidden_states + + +class Encoder(nn.Module): + def __init__( + self, + in_channels=3, + out_channels=3, + block_out_channels=(64,), + layers_per_block=2, + norm_num_groups=32, + double_z=True, + ): + super().__init__() + self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + + self.down_blocks = nn.ModuleList([]) + output_channel = block_out_channels[0] + for i, block_out_channel in enumerate(block_out_channels): + input_channel = output_channel + output_channel = block_out_channel + is_final_block = i == len(block_out_channels) - 1 + + block = DownEncoderBlock2D( + input_channel, + output_channel, + num_layers=layers_per_block, + resnet_groups=norm_num_groups, + add_downsample=not is_final_block, + ) + self.down_blocks.append(block) + + self.mid_block = UNetMidBlock2D( + block_out_channels[-1], + resnet_groups=norm_num_groups, + ) + + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1) + + def forward(self, x): + x = self.conv_in(x) + for block in self.down_blocks: + x = block(x) + x = self.mid_block(x) + x = self.conv_norm_out(x) + x = self.conv_act(x) + x = self.conv_out(x) + return x + + +class Decoder(nn.Module): + def __init__( + self, + in_channels=3, + out_channels=3, + block_out_channels=(64,), + layers_per_block=2, + norm_num_groups=32, + ): + super().__init__() + self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1) + + self.mid_block = UNetMidBlock2D( + block_out_channels[-1], + resnet_groups=norm_num_groups, + ) + + self.up_blocks = nn.ModuleList([]) + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + + for i, block_out_channel in enumerate(reversed_block_out_channels): + input_channel = output_channel + output_channel = block_out_channel + is_final_block = i == len(block_out_channels) - 1 + block = UpDecoderBlock2D( + input_channel, + output_channel, + num_layers=layers_per_block + 1, + resnet_groups=norm_num_groups, + add_upsample=not is_final_block, + ) + self.up_blocks.append(block) + + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1) + + def forward(self, x): + x = self.conv_in(x) + x = self.mid_block(x) + for block in self.up_blocks: + x = block(x) + x = self.conv_norm_out(x) + x = self.conv_act(x) + x = self.conv_out(x) + return x + + +class AutoencoderKL(nn.Module): + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str] = ("DownEncoderBlock2D",), + up_block_types: Tuple[str] = ("UpDecoderBlock2D",), + block_out_channels: Tuple[int] = (64,), + layers_per_block: int = 1, + act_fn: str = "silu", + latent_channels: int = 4, + norm_num_groups: int = 32, + sample_size: int = 32, + scaling_factor: float = 0.18215, + shift_factor: Optional[float] = None, + force_upcast: bool = True, + use_quant_conv: bool = True, + use_post_quant_conv: bool = True, + mid_block_add_attention: bool = True, + **kwargs, + ): + super().__init__() + self.config = AutoencoderConfig( + in_channels=in_channels, + out_channels=out_channels, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + latent_channels=latent_channels, + scaling_factor=scaling_factor, + shift_factor=shift_factor, + ) + + self.encoder = Encoder( + in_channels=in_channels, + out_channels=latent_channels, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + double_z=True, + ) + + self.decoder = Decoder( + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + ) + + self.quant_conv = nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1) if use_quant_conv else None + self.post_quant_conv = nn.Conv2d(latent_channels, latent_channels, 1) if use_post_quant_conv else None + + @property + def dtype(self): + return next(self.parameters()).dtype + + def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput: + if self.post_quant_conv is not None: + z = self.post_quant_conv(z) + + dec = self.decoder(z) + + if not return_dict: + return (dec,) + + return AutoencoderKLOutput(sample=dec) diff --git a/unet/Z-Image/src/zimage/pipeline.py b/unet/Z-Image/src/zimage/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..899673f9d9fb6c082905c7cabdaa0a8f4726f6b8 --- /dev/null +++ b/unet/Z-Image/src/zimage/pipeline.py @@ -0,0 +1,293 @@ +"""Z-Image Pipeline.""" + +import inspect +from typing import List, Optional, Union + +from loguru import logger +import torch + +from config import ( + BASE_IMAGE_SEQ_LEN, + BASE_SHIFT, + DEFAULT_CFG_TRUNCATION, + DEFAULT_GUIDANCE_SCALE, + DEFAULT_HEIGHT, + DEFAULT_INFERENCE_STEPS, + DEFAULT_MAX_SEQUENCE_LENGTH, + DEFAULT_WIDTH, + MAX_IMAGE_SEQ_LEN, + MAX_SHIFT, +) + + +def calculate_shift( + image_seq_len, + base_seq_len: int = BASE_IMAGE_SEQ_LEN, + max_seq_len: int = MAX_IMAGE_SEQ_LEN, + base_shift: float = BASE_SHIFT, + max_shift: float = MAX_SHIFT, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed.") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError(f"The scheduler does not support custom timestep schedules.") + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError(f"The scheduler does not support custom sigmas schedules.") + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +@torch.no_grad() +def generate( + transformer, + vae, + text_encoder, + tokenizer, + scheduler, + prompt: Union[str, List[str]], + height: int = DEFAULT_HEIGHT, + width: int = DEFAULT_WIDTH, + num_inference_steps: int = DEFAULT_INFERENCE_STEPS, + guidance_scale: float = DEFAULT_GUIDANCE_SCALE, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: int = 1, + generator: Optional[torch.Generator] = None, + cfg_normalization: bool = False, + cfg_truncation: float = DEFAULT_CFG_TRUNCATION, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, + output_type: str = "pil", +): + device = next(transformer.parameters()).device + + if hasattr(vae, "config") and hasattr(vae.config, "block_out_channels"): + vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1) + else: + vae_scale_factor = 8 + vae_scale = vae_scale_factor * 2 + + if height % vae_scale != 0: + raise ValueError(f"Height must be divisible by {vae_scale} (got {height}).") + if width % vae_scale != 0: + raise ValueError(f"Width must be divisible by {vae_scale} (got {width}).") + + if isinstance(prompt, str): + batch_size = 1 + prompt = [prompt] + else: + batch_size = len(prompt) + + do_classifier_free_guidance = guidance_scale > 1.0 + logger.info(f"Generating image: {height}x{width}, steps={num_inference_steps}, cfg={guidance_scale}") + + formatted_prompts = [] + for p in prompt: + messages = [{"role": "user", "content": p}] + formatted_prompt = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + formatted_prompts.append(formatted_prompt) + + text_inputs = tokenizer( + formatted_prompts, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device) + prompt_masks = text_inputs.attention_mask.to(device).bool() + + prompt_embeds = text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_masks, + output_hidden_states=True, + ).hidden_states[-2] + + prompt_embeds_list = [] + for i in range(len(prompt_embeds)): + prompt_embeds_list.append(prompt_embeds[i][prompt_masks[i]]) + + negative_prompt_embeds_list = [] + if do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = ["" for _ in prompt] + elif isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] + + neg_formatted = [] + for p in negative_prompt: + messages = [{"role": "user", "content": p}] + formatted_prompt = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + neg_formatted.append(formatted_prompt) + + neg_inputs = tokenizer( + neg_formatted, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + neg_input_ids = neg_inputs.input_ids.to(device) + neg_masks = neg_inputs.attention_mask.to(device).bool() + + neg_embeds = text_encoder( + input_ids=neg_input_ids, + attention_mask=neg_masks, + output_hidden_states=True, + ).hidden_states[-2] + + for i in range(len(neg_embeds)): + negative_prompt_embeds_list.append(neg_embeds[i][neg_masks[i]]) + + if num_images_per_prompt > 1: + prompt_embeds_list = [pe for pe in prompt_embeds_list for _ in range(num_images_per_prompt)] + if do_classifier_free_guidance: + negative_prompt_embeds_list = [ + npe for npe in negative_prompt_embeds_list for _ in range(num_images_per_prompt) + ] + + height_latent = 2 * (int(height) // vae_scale) + width_latent = 2 * (int(width) // vae_scale) + shape = (batch_size * num_images_per_prompt, transformer.in_channels, height_latent, width_latent) + + latents = torch.randn(shape, generator=generator, device=device, dtype=torch.float32) + + actual_batch_size = batch_size * num_images_per_prompt + image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) + + mu = calculate_shift( + image_seq_len, + scheduler.config.get("base_image_seq_len", 256), + scheduler.config.get("max_image_seq_len", 4096), + scheduler.config.get("base_shift", 0.5), + scheduler.config.get("max_shift", 1.15), + ) + scheduler.sigma_min = 0.0 + scheduler_kwargs = {"mu": mu} + timesteps, num_inference_steps = retrieve_timesteps( + scheduler, + num_inference_steps, + device, + sigmas=None, + **scheduler_kwargs, + ) + + logger.info(f"Sampling loop start: {num_inference_steps} steps") + + from tqdm import tqdm + + # Denoising loop with progress bar + for i, t in enumerate(tqdm(timesteps, desc="Denoising", total=len(timesteps))): + # If current t is 0 and it's the last step, skip computation + if t == 0 and i == len(timesteps) - 1: + logger.debug(f"Step {i+1}/{num_inference_steps} | t: {t.item():.2f} | Skipping last step") + continue + + timestep = t.expand(latents.shape[0]) + timestep = (1000 - timestep) / 1000 + t_norm = timestep[0].item() + + current_guidance_scale = guidance_scale + if do_classifier_free_guidance and cfg_truncation is not None and float(cfg_truncation) <= 1: + if t_norm > cfg_truncation: + current_guidance_scale = 0.0 + + apply_cfg = do_classifier_free_guidance and current_guidance_scale > 0 + + if apply_cfg: + latents_typed = latents.to( + transformer.dtype if hasattr(transformer, "dtype") else next(transformer.parameters()).dtype + ) + latent_model_input = latents_typed.repeat(2, 1, 1, 1) + prompt_embeds_model_input = prompt_embeds_list + negative_prompt_embeds_list + timestep_model_input = timestep.repeat(2) + else: + latent_model_input = latents.to(next(transformer.parameters()).dtype) + prompt_embeds_model_input = prompt_embeds_list + timestep_model_input = timestep + + latent_model_input = latent_model_input.unsqueeze(2) + latent_model_input_list = list(latent_model_input.unbind(dim=0)) + + model_out_list = transformer( + latent_model_input_list, + timestep_model_input, + prompt_embeds_model_input, + )[0] + + if apply_cfg: + pos_out = model_out_list[:actual_batch_size] + neg_out = model_out_list[actual_batch_size:] + noise_pred = [] + for j in range(actual_batch_size): + pos = pos_out[j].float() + neg = neg_out[j].float() + pred = pos + current_guidance_scale * (pos - neg) + + if cfg_normalization and float(cfg_normalization) > 0.0: + ori_pos_norm = torch.linalg.vector_norm(pos) + new_pos_norm = torch.linalg.vector_norm(pred) + max_new_norm = ori_pos_norm * float(cfg_normalization) + if new_pos_norm > max_new_norm: + pred = pred * (max_new_norm / new_pos_norm) + noise_pred.append(pred) + noise_pred = torch.stack(noise_pred, dim=0) + else: + noise_pred = torch.stack([t.float() for t in model_out_list], dim=0) + + noise_pred = -noise_pred.squeeze(2) + latents = scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0] + assert latents.dtype == torch.float32 + + if output_type == "latent": + return latents + + shift_factor = getattr(vae.config, "shift_factor", 0.0) or 0.0 + latents = (latents.to(vae.dtype) / vae.config.scaling_factor) + shift_factor + image = vae.decode(latents, return_dict=False)[0] + + if output_type == "pil": + from PIL import Image + + image = (image / 2 + 0.5).clamp(0, 1) + image = image.cpu().permute(0, 2, 3, 1).float().numpy() + image = (image * 255).round().astype("uint8") + image = [Image.fromarray(img) for img in image] + + return image diff --git a/unet/Z-Image/src/zimage/transformer.py b/unet/Z-Image/src/zimage/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..866a00ba95c7df49bbfaf9d2f7bc3fa3bfa9d431 --- /dev/null +++ b/unet/Z-Image/src/zimage/transformer.py @@ -0,0 +1,571 @@ +"""Z-Image Transformer.""" + +import math +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence + +from config import ( + ADALN_EMBED_DIM, + FREQUENCY_EMBEDDING_SIZE, + MAX_PERIOD, + ROPE_AXES_DIMS, + ROPE_AXES_LENS, + ROPE_THETA, + SEQ_MULTI_OF, +) + + +class TimestepEmbedder(nn.Module): + def __init__(self, out_size, mid_size=None, frequency_embedding_size=FREQUENCY_EMBEDDING_SIZE): + super().__init__() + if mid_size is None: + mid_size = out_size + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, mid_size, bias=True), + nn.SiLU(), + nn.Linear(mid_size, out_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=MAX_PERIOD): + with torch.amp.autocast("cuda", enabled=False): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + weight_dtype = self.mlp[0].weight.dtype + if weight_dtype.is_floating_point: + t_freq = t_freq.to(weight_dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return output * self.weight + + +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + with torch.amp.autocast("cuda", enabled=False): + x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x * freqs_cis).flatten(3) + return x_out.type_as(x_in) + + +class ZImageAttention(nn.Module): + _attention_backend = None + + def __init__(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: bool = True, eps: float = 1e-5): + super().__init__() + self.n_heads = n_heads + self.n_kv_heads = n_kv_heads + self.head_dim = dim // n_heads + + self.to_q = nn.Linear(dim, n_heads * self.head_dim, bias=False) + self.to_k = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False) + self.to_v = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False) + self.to_out = nn.ModuleList([nn.Linear(n_heads * self.head_dim, dim, bias=False)]) + + self.norm_q = RMSNorm(self.head_dim, eps=eps) if qk_norm else None + self.norm_k = RMSNorm(self.head_dim, eps=eps) if qk_norm else None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + freqs_cis: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + query = query.unflatten(-1, (self.n_heads, -1)) + key = key.unflatten(-1, (self.n_kv_heads, -1)) + value = value.unflatten(-1, (self.n_kv_heads, -1)) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + if freqs_cis is not None: + query = apply_rotary_emb(query, freqs_cis) + key = apply_rotary_emb(key, freqs_cis) + + dtype = query.dtype + query, key = query.to(dtype), key.to(dtype) + + # Dispatch + from utils.attention import dispatch_attention + + hidden_states = dispatch_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False, backend=self._attention_backend + ) + + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(dtype) + + output = self.to_out[0](hidden_states) + return output + + +class ZImageTransformerBlock(nn.Module): + def __init__( + self, + layer_id: int, + dim: int, + n_heads: int, + n_kv_heads: int, + norm_eps: float, + qk_norm: bool, + modulation=True, + ): + super().__init__() + self.dim = dim + self.head_dim = dim // n_heads + self.layer_id = layer_id + self.modulation = modulation + + self.attention = ZImageAttention(dim, n_heads, n_kv_heads, qk_norm, norm_eps) + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + + self.attention_norm1 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.attention_norm2 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + if modulation: + self.adaLN_modulation = nn.ModuleList([nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True)]) + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + freqs_cis: torch.Tensor, + adaln_input: Optional[torch.Tensor] = None, + ): + if self.modulation: + assert adaln_input is not None + scale_msa, gate_msa, scale_mlp, gate_mlp = ( + self.adaLN_modulation[0](adaln_input).unsqueeze(1).chunk(4, dim=2) + ) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + attn_out = self.attention( + self.attention_norm1(x) * scale_msa, + attention_mask=attn_mask, + freqs_cis=freqs_cis, + ) + x = x + gate_msa * self.attention_norm2(attn_out) + x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp)) + else: + attn_out = self.attention( + self.attention_norm1(x), + attention_mask=attn_mask, + freqs_cis=freqs_cis, + ) + x = x + self.attention_norm2(attn_out) + x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x))) + + return x + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True), + ) + + def forward(self, x, c): + scale = 1.0 + self.adaLN_modulation(c) + x = self.norm_final(x) * scale.unsqueeze(1) + x = self.linear(x) + return x + + +class RopeEmbedder: + def __init__( + self, + theta: float = ROPE_THETA, + axes_dims: List[int] = ROPE_AXES_DIMS, + axes_lens: List[int] = ROPE_AXES_LENS, + ): + self.theta = theta + self.axes_dims = axes_dims + self.axes_lens = axes_lens + assert len(axes_dims) == len(axes_lens) + self.freqs_cis = None + + @staticmethod + def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = ROPE_THETA): + with torch.device("cpu"): + freqs_cis = [] + for i, (d, e) in enumerate(zip(dim, end)): + freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)) + timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) + freqs = torch.outer(timestep, freqs).float() + freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64) + freqs_cis.append(freqs_cis_i) + return freqs_cis + + def __call__(self, ids: torch.Tensor): + assert ids.ndim == 2 + assert ids.shape[-1] == len(self.axes_dims) + device = ids.device + + if self.freqs_cis is None: + self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta) + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + else: + if self.freqs_cis[0].device != device: + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + + result = [] + for i in range(len(self.axes_dims)): + index = ids[:, i] + result.append(self.freqs_cis[i][index]) + return torch.cat(result, dim=-1) + + +class ZImageTransformer2DModel(nn.Module): + def __init__( + self, + all_patch_size=(2,), + all_f_patch_size=(1,), + in_channels=16, + dim=3840, + n_layers=30, + n_refiner_layers=2, + n_heads=30, + n_kv_heads=30, + norm_eps=1e-5, + qk_norm=True, + cap_feat_dim=2560, + rope_theta=ROPE_THETA, + t_scale=1000.0, + axes_dims=ROPE_AXES_DIMS, + axes_lens=ROPE_AXES_LENS, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.all_patch_size = all_patch_size + self.all_f_patch_size = all_f_patch_size + self.dim = dim + self.n_heads = n_heads + self.rope_theta = rope_theta + self.t_scale = t_scale + + assert len(all_patch_size) == len(all_f_patch_size) + + all_x_embedder = {} + all_final_layer = {} + for patch_size, f_patch_size in zip(all_patch_size, all_f_patch_size): + x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True) + all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder + final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels) + all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer + + self.all_x_embedder = nn.ModuleDict(all_x_embedder) + self.all_final_layer = nn.ModuleDict(all_final_layer) + + self.noise_refiner = nn.ModuleList( + [ + ZImageTransformerBlock(1000 + layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=True) + for layer_id in range(n_refiner_layers) + ] + ) + + self.context_refiner = nn.ModuleList( + [ + ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=False) + for layer_id in range(n_refiner_layers) + ] + ) + + self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024) + self.cap_embedder = nn.Sequential( + RMSNorm(cap_feat_dim, eps=norm_eps), + nn.Linear(cap_feat_dim, dim, bias=True), + ) + + self.x_pad_token = nn.Parameter(torch.empty((1, dim))) + self.cap_pad_token = nn.Parameter(torch.empty((1, dim))) + + self.layers = nn.ModuleList( + [ + ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm) + for layer_id in range(n_layers) + ] + ) + + head_dim = dim // n_heads + assert head_dim == sum(axes_dims) + self.axes_dims = axes_dims + self.axes_lens = axes_lens + + self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens) + + def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size) -> List[torch.Tensor]: + pH = pW = patch_size + pF = f_patch_size + bsz = len(x) + assert len(size) == bsz + for i in range(bsz): + F, H, W = size[i] + ori_len = (F // pF) * (H // pH) * (W // pW) + x[i] = ( + x[i][:ori_len] + .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels) + .permute(6, 0, 3, 1, 4, 2, 5) + .reshape(self.out_channels, F, H, W) + ) + return x + + @staticmethod + def create_coordinate_grid(size, start=None, device=None): + if start is None: + start = (0 for _ in size) + axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)] + grids = torch.meshgrid(axes, indexing="ij") + return torch.stack(grids, dim=-1) + + def patchify_and_embed( + self, + all_image: List[torch.Tensor], + all_cap_feats: List[torch.Tensor], + patch_size: int, + f_patch_size: int, + ): + pH = pW = patch_size + pF = f_patch_size + device = all_image[0].device + + all_image_out = [] + all_image_size = [] + all_image_pos_ids = [] + all_image_pad_mask = [] + all_cap_pos_ids = [] + all_cap_pad_mask = [] + all_cap_feats_out = [] + + for _, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)): + cap_ori_len = len(cap_feat) + cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF + cap_padded_pos_ids = self.create_coordinate_grid( + size=(cap_ori_len + cap_padding_len, 1, 1), + start=(1, 0, 0), + device=device, + ).flatten(0, 2) + all_cap_pos_ids.append(cap_padded_pos_ids) + # pad mask + all_cap_pad_mask.append( + torch.cat( + [ + torch.zeros((cap_ori_len,), dtype=torch.bool, device=device), + torch.ones((cap_padding_len,), dtype=torch.bool, device=device), + ], + dim=0, + ) + if cap_padding_len > 0 + else torch.zeros((cap_ori_len,), dtype=torch.bool, device=device) + ) + # padded feature + all_cap_feats_out.append( + torch.cat( + [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], + dim=0, + ) + if cap_padding_len > 0 + else cap_feat + ) + + C, F, H, W = image.size() + all_image_size.append((F, H, W)) + F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + + image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) + image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C) + + image_ori_len = len(image) + image_padding_len = (-image_ori_len) % SEQ_MULTI_OF + + image_ori_pos_ids = self.create_coordinate_grid( + size=(F_tokens, H_tokens, W_tokens), + start=(cap_ori_len + cap_padding_len + 1, 0, 0), + device=device, + ).flatten(0, 2) + image_padded_pos_ids = torch.cat( + [ + image_ori_pos_ids, + self.create_coordinate_grid(size=(1, 1, 1), start=(0, 0, 0), device=device) + .flatten(0, 2) + .repeat(image_padding_len, 1), + ], + dim=0, + ) + all_image_pos_ids.append(image_padded_pos_ids if image_padding_len > 0 else image_ori_pos_ids) + # pad mask + image_pad_mask = torch.cat( + [ + torch.zeros((image_ori_len,), dtype=torch.bool, device=device), + torch.ones((image_padding_len,), dtype=torch.bool, device=device), + ], + dim=0, + ) + all_image_pad_mask.append( + image_pad_mask + if image_padding_len > 0 + else torch.zeros((image_ori_len,), dtype=torch.bool, device=device) + ) + # padded feature + image_padded_feat = torch.cat( + [image, image[-1:].repeat(image_padding_len, 1)], + dim=0, + ) + all_image_out.append(image_padded_feat if image_padding_len > 0 else image) + + return ( + all_image_out, + all_cap_feats_out, + all_image_size, + all_image_pos_ids, + all_cap_pos_ids, + all_image_pad_mask, + all_cap_pad_mask, + ) + + def forward( + self, + x: List[torch.Tensor], + t, + cap_feats: List[torch.Tensor], + patch_size=2, + f_patch_size=1, + ): + assert patch_size in self.all_patch_size + assert f_patch_size in self.all_f_patch_size + + bsz = len(x) + device = x[0].device + t = t * self.t_scale + t = self.t_embedder(t) + + ( + x, + cap_feats, + x_size, + x_pos_ids, + cap_pos_ids, + x_inner_pad_mask, + cap_inner_pad_mask, + ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) + + x_item_seqlens = [len(_) for _ in x] + assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) + x_max_item_seqlen = max(x_item_seqlens) + + x = torch.cat(x, dim=0) + x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) + + adaln_input = t.type_as(x) + x[torch.cat(x_inner_pad_mask)] = self.x_pad_token + x = list(x.split(x_item_seqlens, dim=0)) + x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split([len(_) for _ in x_pos_ids], dim=0)) + + x = pad_sequence(x, batch_first=True, padding_value=0.0) + x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) + # Clarify the length matches to satisfy Dynamo due to "Symbolic Shape Inference" to avoid compilation errors + x_freqs_cis = x_freqs_cis[:, : x.shape[1]] + + x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(x_item_seqlens): + x_attn_mask[i, :seq_len] = 1 + + for layer in self.noise_refiner: + x = layer(x, x_attn_mask, x_freqs_cis, adaln_input) + + cap_item_seqlens = [len(_) for _ in cap_feats] + assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) + cap_max_item_seqlen = max(cap_item_seqlens) + + cap_feats = torch.cat(cap_feats, dim=0) + cap_feats = self.cap_embedder(cap_feats) + cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token + cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) + cap_freqs_cis = list( + self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split([len(_) for _ in cap_pos_ids], dim=0) + ) + + cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0) + cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) + cap_freqs_cis = cap_freqs_cis[:, : cap_feats.shape[1]] # same for dynamo compatibility + + cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(cap_item_seqlens): + cap_attn_mask[i, :seq_len] = 1 + + for layer in self.context_refiner: + cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis) + + unified = [] + unified_freqs_cis = [] + for i in range(bsz): + x_len = x_item_seqlens[i] + cap_len = cap_item_seqlens[i] + unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) + unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]])) + unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] + assert unified_item_seqlens == [len(_) for _ in unified] + unified_max_item_seqlen = max(unified_item_seqlens) + + unified = pad_sequence(unified, batch_first=True, padding_value=0.0) + unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) + unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(unified_item_seqlens): + unified_attn_mask[i, :seq_len] = 1 + + for layer in self.layers: + unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input) + + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input) + unified = list(unified.unbind(dim=0)) + x = self.unpatchify(unified, x_size, patch_size, f_patch_size) + + return x, {} diff --git a/upscale_models/1x-ITF-SkinDiffDetail-Lite-v1.pth b/upscale_models/1x-ITF-SkinDiffDetail-Lite-v1.pth new file mode 100644 index 0000000000000000000000000000000000000000..dc8dc4c411065da9affd4585c150e5d72d214966 --- /dev/null +++ b/upscale_models/1x-ITF-SkinDiffDetail-Lite-v1.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:94d368b633614958f84f335b129fd85abd30200e8fbc575b859ba6762116222b +size 20099337 diff --git a/upscale_models/1x_PureVision.pth b/upscale_models/1x_PureVision.pth new file mode 100644 index 0000000000000000000000000000000000000000..bb9ac38f74b23980aee414cc5c18a9e54d3396bb --- /dev/null +++ b/upscale_models/1x_PureVision.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3c3bec2111d0f7c3926f3171f37ce28e72502b744e084c566d8960c06a6a06a3 +size 67120607 diff --git a/upscale_models/2x_PureVision.pth b/upscale_models/2x_PureVision.pth new file mode 100644 index 0000000000000000000000000000000000000000..9e5ab70374b64cfe28f4e4f48515624695d04570 --- /dev/null +++ b/upscale_models/2x_PureVision.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c49109c8257b80b3cddf0110786ad11c1ca05214a470bfcc7fd49b6461dfcaee +size 67037663 diff --git a/upscale_models/4x-ClearRealityV1.pth b/upscale_models/4x-ClearRealityV1.pth new file mode 100644 index 0000000000000000000000000000000000000000..417bc92902a2b343e3f385008255d213362fe197 --- /dev/null +++ b/upscale_models/4x-ClearRealityV1.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a4cd3a25b00e0be949d4302fc774eb4d7f2ed5f47cdb51551e2d75fa6562e51e +size 9016074 diff --git a/upscale_models/4x-UltraSharp.pth b/upscale_models/4x-UltraSharp.pth new file mode 100644 index 0000000000000000000000000000000000000000..9f3bb839bebd6cd26c94122b7651261d0b346a50 --- /dev/null +++ b/upscale_models/4x-UltraSharp.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a5812231fc936b42af08a5edba784195495d303d5b3248c24489ef0c4021fe01 +size 66961958 diff --git a/upscale_models/4xFFHQDAT.safetensors b/upscale_models/4xFFHQDAT.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f20079723562ec68ba513b26e0b62bb42ec28ddc --- /dev/null +++ b/upscale_models/4xFFHQDAT.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fa91faa8c1f72c32646d71abf51e952c81b4984948e89e6e5a8c40822a6cf3cc +size 154152604 diff --git a/upscale_models/4xNomos8k_atd_jpg.pth b/upscale_models/4xNomos8k_atd_jpg.pth new file mode 100644 index 0000000000000000000000000000000000000000..5c9cfa7bd4a7516d391195b37a9b48d3f7be6340 --- /dev/null +++ b/upscale_models/4xNomos8k_atd_jpg.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6b2bfb0e822c79594288dd43efaec213b6f0244384bd98db75072b0ce5a729fe +size 81959074 diff --git a/upscale_models/4xNomos8k_span_otf_weak.pth b/upscale_models/4xNomos8k_span_otf_weak.pth new file mode 100644 index 0000000000000000000000000000000000000000..d3d845bef7a3e4a366dfed2e5de8fe07fba641ca --- /dev/null +++ b/upscale_models/4xNomos8k_span_otf_weak.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:26e7ef6483faf93b47f48af262ed4bb8dededa20af34a923e344cb63cafeec0a +size 9015866 diff --git a/upscale_models/4x_NMKD-Siax_200k.pth b/upscale_models/4x_NMKD-Siax_200k.pth new file mode 100644 index 0000000000000000000000000000000000000000..3e90a2b5e1bb5921fb8e16e2d9ed1983af75d963 --- /dev/null +++ b/upscale_models/4x_NMKD-Siax_200k.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:560424d9f68625713fc47e9e7289a98aabe1d744e1cd6a9ae5a35e9957fd127e +size 66957746 diff --git a/upscale_models/4x_NMKD-Superscale-SP_178000_G.pth b/upscale_models/4x_NMKD-Superscale-SP_178000_G.pth new file mode 100644 index 0000000000000000000000000000000000000000..0216b8a91fcd58d8e3fa99a77383a331eb20b615 --- /dev/null +++ b/upscale_models/4x_NMKD-Superscale-SP_178000_G.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d1b0078fe71446e0469d8d4df59e96baa80d83cda600d68237d655830821bcc +size 66958607 diff --git a/upscale_models/4x_foolhardy_Remacri.pth b/upscale_models/4x_foolhardy_Remacri.pth new file mode 100644 index 0000000000000000000000000000000000000000..7dc76f8e1e196da607cabbbf9f58fa4acea73930 --- /dev/null +++ b/upscale_models/4x_foolhardy_Remacri.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e1a73bd89c2da1ae494774746398689048b5a892bd9653e146713f9df8bca86a +size 67025055 diff --git a/upscale_models/RealESRGAN_x4plus.pth b/upscale_models/RealESRGAN_x4plus.pth new file mode 100644 index 0000000000000000000000000000000000000000..9ddced536d07803300536317fef662bb499bca71 --- /dev/null +++ b/upscale_models/RealESRGAN_x4plus.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4fa0d38905f75ac06eb49a7951b426670021be3018265fd191d2125df9d682f1 +size 67040989 diff --git a/vae_approx/taew2_1.pth b/vae_approx/taew2_1.pth new file mode 100644 index 0000000000000000000000000000000000000000..5d391ecb617c387bc3fdc17aa4d7daf8220ae522 --- /dev/null +++ b/vae_approx/taew2_1.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e +size 22678901 diff --git a/vitmatte/model.safetensors b/vitmatte/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..de17d1740149f189b7f67fc93db6eef3e8a48e84 --- /dev/null +++ b/vitmatte/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bda9289db1bb6762d978b42d1c62ae3f34daf7497171a347a1d09657efd788cb +size 103294572 diff --git a/vitmatte/pytorch_model.bin b/vitmatte/pytorch_model.bin new file mode 100644 index 0000000000000000000000000000000000000000..2966b406de9a9a4a09d7807e7c031a4e6d300cd5 --- /dev/null +++ b/vitmatte/pytorch_model.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6ec6aed44bc8d8ab7f4d0ff46da3520a534cf5a97a8262404ff6efa9ae33b1e5 +size 103349013