Upload W-Flow-diffusers checkpoints

.gitattributes

@@ -33,3 +33,4 @@ 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
+demo.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.py[cod]
+.DS_Store

README.md

@@ -0,0 +1,107 @@
+---
+license: apache-2.0
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+- diffusers
+- w-flow
+- image-generation
+- class-conditional
+- one-step
+- imagenet
+widget:
+- text: golden retriever
+  output:
+    url: demo.png
+language:
+- en
+---
+
+# W-Flow-diffusers
+
+Native [diffusers](https://github.com/huggingface/diffusers) checkpoints for **W-Flow** ([One-Step Generative Modeling via Wasserstein Gradient Flows](https://arxiv.org/abs/2605.11755)). Code: [hanjq17/W-Flow](https://github.com/hanjq17/W-Flow). Each variant folder is self-contained:
+
+- `pipeline.py` — `WFlowPipeline`
+- `scheduler/scheduling_wflow.py` — `WFlowScheduler` (one evaluation; the transformer predicts the clean latent)
+- `transformer/transformer_wflow.py` — `WFlowTransformer2DModel`
+- `vae/` — `stabilityai/sd-vae-ft-mse` (`AutoencoderKL`, scaling factor 0.18215)
+
+Load with PyPI `diffusers` and `trust_remote_code=True`. No extra W-Flow package is required. Sampling is one forward pass at 256×256. `guidance_scale` is the official CFG value, which the paper writes as \(w+1\). Weights are the released EMA. The paper reports ImageNet-256 FID 1.52 for W-Flow-B, 1.35 for W-Flow-L, and 1.29 for W-Flow-XL.
+
+## Available checkpoints
+
+| Checkpoint | Path | Resolution | Steps | Guidance |
+| --- | --- | --- | ---: | ---: |
+| W-Flow-B | [`W-Flow-B/`](W-Flow-B/) | 256×256 | 1 | 1.19 |
+| W-Flow-L | [`W-Flow-L/`](W-Flow-L/) | 256×256 | 1 | 1.14 |
+| W-Flow-XL | [`W-Flow-XL/`](W-Flow-XL/) | 256×256 | 1 | 1.09 |
+
+## Demo
+
+![W-Flow-B demo](demo.png)
+
+Class: *golden retriever* (ImageNet 207) — 256×256, 1 step, `guidance_scale=1.19`, `torch_dtype=bfloat16`, seed 1.
+
+## ImageNet class labels
+
+Each variant keeps an English `id2label` map in `model_index.json`.
+
+- `pipe.id2label` — id → English label
+- `pipe.labels` — synonym → id
+- `pipe.get_label_ids("golden retriever")`
+- `pipe(class_labels="golden retriever", ...)` — string labels resolve automatically
+
+`height` and `width` default to 256. The patch embedding is fixed to a 32×32 latent, so other sizes are rejected. On CUDA the generator autocasts to bfloat16.
+
+## Inference
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/W-Flow-diffusers",
+    subfolder="W-Flow-B",
+    custom_pipeline="pipeline.py",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+print(pipe.id2label[207])
+image = pipe(
+    class_labels="golden retriever",
+    num_inference_steps=1,
+    guidance_scale=1.19,
+    generator=torch.Generator(device="cuda").manual_seed(1),
+).images[0]
+image.save("demo.png")
+```
+
+Local checkout:
+
+```python
+from pathlib import Path
+import torch
+from diffusers import DiffusionPipeline
+
+model_dir = Path("W-Flow-diffusers/W-Flow-B")
+pipe = DiffusionPipeline.from_pretrained(
+    str(model_dir),
+    local_files_only=True,
+    custom_pipeline=str(model_dir / "pipeline.py"),
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+```
+
+## Citation
+
+```bibtex
+@article{han2026one,
+  title   = {One-Step Generative Modeling via Wasserstein Gradient Flows},
+  author  = {Han, Jiaqi and Li, Puheng and Guo, Qiushan and Xu, Renyuan and Ermon, Stefano and Cand{\`e}s, Emmanuel J.},
+  journal = {arXiv preprint arXiv:2605.11755},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2605.11755}
+}
+```

W-Flow-B/model_index.json

@@ -0,0 +1,1021 @@
+{
+  "_class_name": [
+    "pipeline",
+    "WFlowPipeline"
+  ],
+  "_diffusers_version": "0.39.0",
+  "scheduler": [
+    "scheduling_wflow",
+    "WFlowScheduler"
+  ],
+  "transformer": [
+    "transformer_wflow",
+    "WFlowTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ],
+  "id2label": {
+    "0": "tench, Tinca tinca",
+    "1": "goldfish, Carassius auratus",
+    "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
+    "3": "tiger shark, Galeocerdo cuvieri",
+    "4": "hammerhead, hammerhead shark",
+    "5": "electric ray, crampfish, numbfish, torpedo",
+    "6": "stingray",
+    "7": "cock",
+    "8": "hen",
+    "9": "ostrich, Struthio camelus",
+    "10": "brambling, Fringilla montifringilla",
+    "11": "goldfinch, Carduelis carduelis",
+    "12": "house finch, linnet, Carpodacus mexicanus",
+    "13": "junco, snowbird",
+    "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+    "15": "robin, American robin, Turdus migratorius",
+    "16": "bulbul",
+    "17": "jay",
+    "18": "magpie",
+    "19": "chickadee",
+    "20": "water ouzel, dipper",
+    "21": "kite",
+    "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
+    "23": "vulture",
+    "24": "great grey owl, great gray owl, Strix nebulosa",
+    "25": "European fire salamander, Salamandra salamandra",
+    "26": "common newt, Triturus vulgaris",
+    "27": "eft",
+    "28": "spotted salamander, Ambystoma maculatum",
+    "29": "axolotl, mud puppy, Ambystoma mexicanum",
+    "30": "bullfrog, Rana catesbeiana",
+    "31": "tree frog, tree-frog",
+    "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+    "33": "loggerhead, loggerhead turtle, Caretta caretta",
+    "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
+    "35": "mud turtle",
+    "36": "terrapin",
+    "37": "box turtle, box tortoise",
+    "38": "banded gecko",
+    "39": "common iguana, iguana, Iguana iguana",
+    "40": "American chameleon, anole, Anolis carolinensis",
+    "41": "whiptail, whiptail lizard",
+    "42": "agama",
+    "43": "frilled lizard, Chlamydosaurus kingi",
+    "44": "alligator lizard",
+    "45": "Gila monster, Heloderma suspectum",
+    "46": "green lizard, Lacerta viridis",
+    "47": "African chameleon, Chamaeleo chamaeleon",
+    "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
+    "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
+    "50": "American alligator, Alligator mississipiensis",
+    "51": "triceratops",
+    "52": "thunder snake, worm snake, Carphophis amoenus",
+    "53": "ringneck snake, ring-necked snake, ring snake",
+    "54": "hognose snake, puff adder, sand viper",
+    "55": "green snake, grass snake",
+    "56": "king snake, kingsnake",
+    "57": "garter snake, grass snake",
+    "58": "water snake",
+    "59": "vine snake",
+    "60": "night snake, Hypsiglena torquata",
+    "61": "boa constrictor, Constrictor constrictor",
+    "62": "rock python, rock snake, Python sebae",
+    "63": "Indian cobra, Naja naja",
+    "64": "green mamba",
+    "65": "sea snake",
+    "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+    "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+    "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
+    "69": "trilobite",
+    "70": "harvestman, daddy longlegs, Phalangium opilio",
+    "71": "scorpion",
+    "72": "black and gold garden spider, Argiope aurantia",
+    "73": "barn spider, Araneus cavaticus",
+    "74": "garden spider, Aranea diademata",
+    "75": "black widow, Latrodectus mactans",
+    "76": "tarantula",
+    "77": "wolf spider, hunting spider",
+    "78": "tick",
+    "79": "centipede",
+    "80": "black grouse",
+    "81": "ptarmigan",
+    "82": "ruffed grouse, partridge, Bonasa umbellus",
+    "83": "prairie chicken, prairie grouse, prairie fowl",
+    "84": "peacock",
+    "85": "quail",
+    "86": "partridge",
+    "87": "African grey, African gray, Psittacus erithacus",
+    "88": "macaw",
+    "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+    "90": "lorikeet",
+    "91": "coucal",
+    "92": "bee eater",
+    "93": "hornbill",
+    "94": "hummingbird",
+    "95": "jacamar",
+    "96": "toucan",
+    "97": "drake",
+    "98": "red-breasted merganser, Mergus serrator",
+    "99": "goose",
+    "100": "black swan, Cygnus atratus",
+    "101": "tusker",
+    "102": "echidna, spiny anteater, anteater",
+    "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",
+    "104": "wallaby, brush kangaroo",
+    "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",
+    "106": "wombat",
+    "107": "jellyfish",
+    "108": "sea anemone, anemone",
+    "109": "brain coral",
+    "110": "flatworm, platyhelminth",
+    "111": "nematode, nematode worm, roundworm",
+    "112": "conch",
+    "113": "snail",
+    "114": "slug",
+    "115": "sea slug, nudibranch",
+    "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore",
+    "117": "chambered nautilus, pearly nautilus, nautilus",
+    "118": "Dungeness crab, Cancer magister",
+    "119": "rock crab, Cancer irroratus",
+    "120": "fiddler crab",
+    "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica",
+    "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus",
+    "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish",
+    "124": "crayfish, crawfish, crawdad, crawdaddy",
+    "125": "hermit crab",
+    "126": "isopod",
+    "127": "white stork, Ciconia ciconia",
+    "128": "black stork, Ciconia nigra",
+    "129": "spoonbill",
+    "130": "flamingo",
+    "131": "little blue heron, Egretta caerulea",
+    "132": "American egret, great white heron, Egretta albus",
+    "133": "bittern",
+    "134": "crane",
+    "135": "limpkin, Aramus pictus",
+    "136": "European gallinule, Porphyrio porphyrio",
+    "137": "American coot, marsh hen, mud hen, water hen, Fulica americana",
+    "138": "bustard",
+    "139": "ruddy turnstone, Arenaria interpres",
+    "140": "red-backed sandpiper, dunlin, Erolia alpina",
+    "141": "redshank, Tringa totanus",
+    "142": "dowitcher",
+    "143": "oystercatcher, oyster catcher",
+    "144": "pelican",
+    "145": "king penguin, Aptenodytes patagonica",
+    "146": "albatross, mollymawk",
+    "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus",
+    "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca",
+    "149": "dugong, Dugong dugon",
+    "150": "sea lion",
+    "151": "Chihuahua",
+    "152": "Japanese spaniel",
+    "153": "Maltese dog, Maltese terrier, Maltese",
+    "154": "Pekinese, Pekingese, Peke",
+    "155": "Shih-Tzu",
+    "156": "Blenheim spaniel",
+    "157": "papillon",
+    "158": "toy terrier",
+    "159": "Rhodesian ridgeback",
+    "160": "Afghan hound, Afghan",
+    "161": "basset, basset hound",
+    "162": "beagle",
+    "163": "bloodhound, sleuthhound",
+    "164": "bluetick",
+    "165": "black-and-tan coonhound",
+    "166": "Walker hound, Walker foxhound",
+    "167": "English foxhound",
+    "168": "redbone",
+    "169": "borzoi, Russian wolfhound",
+    "170": "Irish wolfhound",
+    "171": "Italian greyhound",
+    "172": "whippet",
+    "173": "Ibizan hound, Ibizan Podenco",
+    "174": "Norwegian elkhound, elkhound",
+    "175": "otterhound, otter hound",
+    "176": "Saluki, gazelle hound",
+    "177": "Scottish deerhound, deerhound",
+    "178": "Weimaraner",
+    "179": "Staffordshire bullterrier, Staffordshire bull terrier",
+    "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",
+    "181": "Bedlington terrier",
+    "182": "Border terrier",
+    "183": "Kerry blue terrier",
+    "184": "Irish terrier",
+    "185": "Norfolk terrier",
+    "186": "Norwich terrier",
+    "187": "Yorkshire terrier",
+    "188": "wire-haired fox terrier",
+    "189": "Lakeland terrier",
+    "190": "Sealyham terrier, Sealyham",
+    "191": "Airedale, Airedale terrier",
+    "192": "cairn, cairn terrier",
+    "193": "Australian terrier",
+    "194": "Dandie Dinmont, Dandie Dinmont terrier",
+    "195": "Boston bull, Boston terrier",
+    "196": "miniature schnauzer",
+    "197": "giant schnauzer",
+    "198": "standard schnauzer",
+    "199": "Scotch terrier, Scottish terrier, Scottie",
+    "200": "Tibetan terrier, chrysanthemum dog",
+    "201": "silky terrier, Sydney silky",
+    "202": "soft-coated wheaten terrier",
+    "203": "West Highland white terrier",
+    "204": "Lhasa, Lhasa apso",
+    "205": "flat-coated retriever",
+    "206": "curly-coated retriever",
+    "207": "golden retriever",
+    "208": "Labrador retriever",
+    "209": "Chesapeake Bay retriever",
+    "210": "German short-haired pointer",
+    "211": "vizsla, Hungarian pointer",
+    "212": "English setter",
+    "213": "Irish setter, red setter",
+    "214": "Gordon setter",
+    "215": "Brittany spaniel",
+    "216": "clumber, clumber spaniel",
+    "217": "English springer, English springer spaniel",
+    "218": "Welsh springer spaniel",
+    "219": "cocker spaniel, English cocker spaniel, cocker",
+    "220": "Sussex spaniel",
+    "221": "Irish water spaniel",
+    "222": "kuvasz",
+    "223": "schipperke",
+    "224": "groenendael",
+    "225": "malinois",
+    "226": "briard",
+    "227": "kelpie",
+    "228": "komondor",
+    "229": "Old English sheepdog, bobtail",
+    "230": "Shetland sheepdog, Shetland sheep dog, Shetland",
+    "231": "collie",
+    "232": "Border collie",
+    "233": "Bouvier des Flandres, Bouviers des Flandres",
+    "234": "Rottweiler",
+    "235": "German shepherd, German shepherd dog, German police dog, alsatian",
+    "236": "Doberman, Doberman pinscher",
+    "237": "miniature pinscher",
+    "238": "Greater Swiss Mountain dog",
+    "239": "Bernese mountain dog",
+    "240": "Appenzeller",
+    "241": "EntleBucher",
+    "242": "boxer",
+    "243": "bull mastiff",
+    "244": "Tibetan mastiff",
+    "245": "French bulldog",
+    "246": "Great Dane",
+    "247": "Saint Bernard, St Bernard",
+    "248": "Eskimo dog, husky",
+    "249": "malamute, malemute, Alaskan malamute",
+    "250": "Siberian husky",
+    "251": "dalmatian, coach dog, carriage dog",
+    "252": "affenpinscher, monkey pinscher, monkey dog",
+    "253": "basenji",
+    "254": "pug, pug-dog",
+    "255": "Leonberg",
+    "256": "Newfoundland, Newfoundland dog",
+    "257": "Great Pyrenees",
+    "258": "Samoyed, Samoyede",
+    "259": "Pomeranian",
+    "260": "chow, chow chow",
+    "261": "keeshond",
+    "262": "Brabancon griffon",
+    "263": "Pembroke, Pembroke Welsh corgi",
+    "264": "Cardigan, Cardigan Welsh corgi",
+    "265": "toy poodle",
+    "266": "miniature poodle",
+    "267": "standard poodle",
+    "268": "Mexican hairless",
+    "269": "timber wolf, grey wolf, gray wolf, Canis lupus",
+    "270": "white wolf, Arctic wolf, Canis lupus tundrarum",
+    "271": "red wolf, maned wolf, Canis rufus, Canis niger",
+    "272": "coyote, prairie wolf, brush wolf, Canis latrans",
+    "273": "dingo, warrigal, warragal, Canis dingo",
+    "274": "dhole, Cuon alpinus",
+    "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
+    "276": "hyena, hyaena",
+    "277": "red fox, Vulpes vulpes",
+    "278": "kit fox, Vulpes macrotis",
+    "279": "Arctic fox, white fox, Alopex lagopus",
+    "280": "grey fox, gray fox, Urocyon cinereoargenteus",
+    "281": "tabby, tabby cat",
+    "282": "tiger cat",
+    "283": "Persian cat",
+    "284": "Siamese cat, Siamese",
+    "285": "Egyptian cat",
+    "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor",
+    "287": "lynx, catamount",
+    "288": "leopard, Panthera pardus",
+    "289": "snow leopard, ounce, Panthera uncia",
+    "290": "jaguar, panther, Panthera onca, Felis onca",
+    "291": "lion, king of beasts, Panthera leo",
+    "292": "tiger, Panthera tigris",
+    "293": "cheetah, chetah, Acinonyx jubatus",
+    "294": "brown bear, bruin, Ursus arctos",
+    "295": "American black bear, black bear, Ursus americanus, Euarctos americanus",
+    "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus",
+    "297": "sloth bear, Melursus ursinus, Ursus ursinus",
+    "298": "mongoose",
+    "299": "meerkat, mierkat",
+    "300": "tiger beetle",
+    "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle",
+    "302": "ground beetle, carabid beetle",
+    "303": "long-horned beetle, longicorn, longicorn beetle",
+    "304": "leaf beetle, chrysomelid",
+    "305": "dung beetle",
+    "306": "rhinoceros beetle",
+    "307": "weevil",
+    "308": "fly",
+    "309": "bee",
+    "310": "ant, emmet, pismire",
+    "311": "grasshopper, hopper",
+    "312": "cricket",
+    "313": "walking stick, walkingstick, stick insect",
+    "314": "cockroach, roach",
+    "315": "mantis, mantid",
+    "316": "cicada, cicala",
+    "317": "leafhopper",
+    "318": "lacewing, lacewing fly",
+    "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",
+    "320": "damselfly",
+    "321": "admiral",
+    "322": "ringlet, ringlet butterfly",
+    "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus",
+    "324": "cabbage butterfly",
+    "325": "sulphur butterfly, sulfur butterfly",
+    "326": "lycaenid, lycaenid butterfly",
+    "327": "starfish, sea star",
+    "328": "sea urchin",
+    "329": "sea cucumber, holothurian",
+    "330": "wood rabbit, cottontail, cottontail rabbit",
+    "331": "hare",
+    "332": "Angora, Angora rabbit",
+    "333": "hamster",
+    "334": "porcupine, hedgehog",
+    "335": "fox squirrel, eastern fox squirrel, Sciurus niger",
+    "336": "marmot",
+    "337": "beaver",
+    "338": "guinea pig, Cavia cobaya",
+    "339": "sorrel",
+    "340": "zebra",
+    "341": "hog, pig, grunter, squealer, Sus scrofa",
+    "342": "wild boar, boar, Sus scrofa",
+    "343": "warthog",
+    "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius",
+    "345": "ox",
+    "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis",
+    "347": "bison",
+    "348": "ram, tup",
+    "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis",
+    "350": "ibex, Capra ibex",
+    "351": "hartebeest",
+    "352": "impala, Aepyceros melampus",
+    "353": "gazelle",
+    "354": "Arabian camel, dromedary, Camelus dromedarius",
+    "355": "llama",
+    "356": "weasel",
+    "357": "mink",
+    "358": "polecat, fitch, foulmart, foumart, Mustela putorius",
+    "359": "black-footed ferret, ferret, Mustela nigripes",
+    "360": "otter",
+    "361": "skunk, polecat, wood pussy",
+    "362": "badger",
+    "363": "armadillo",
+    "364": "three-toed sloth, ai, Bradypus tridactylus",
+    "365": "orangutan, orang, orangutang, Pongo pygmaeus",
+    "366": "gorilla, Gorilla gorilla",
+    "367": "chimpanzee, chimp, Pan troglodytes",
+    "368": "gibbon, Hylobates lar",
+    "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus",
+    "370": "guenon, guenon monkey",
+    "371": "patas, hussar monkey, Erythrocebus patas",
+    "372": "baboon",
+    "373": "macaque",
+    "374": "langur",
+    "375": "colobus, colobus monkey",
+    "376": "proboscis monkey, Nasalis larvatus",
+    "377": "marmoset",
+    "378": "capuchin, ringtail, Cebus capucinus",
+    "379": "howler monkey, howler",
+    "380": "titi, titi monkey",
+    "381": "spider monkey, Ateles geoffroyi",
+    "382": "squirrel monkey, Saimiri sciureus",
+    "383": "Madagascar cat, ring-tailed lemur, Lemur catta",
+    "384": "indri, indris, Indri indri, Indri brevicaudatus",
+    "385": "Indian elephant, Elephas maximus",
+    "386": "African elephant, Loxodonta africana",
+    "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens",
+    "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca",
+    "389": "barracouta, snoek",
+    "390": "eel",
+    "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch",
+    "392": "rock beauty, Holocanthus tricolor",
+    "393": "anemone fish",
+    "394": "sturgeon",
+    "395": "gar, garfish, garpike, billfish, Lepisosteus osseus",
+    "396": "lionfish",
+    "397": "puffer, pufferfish, blowfish, globefish",
+    "398": "abacus",
+    "399": "abaya",
+    "400": "academic gown, academic robe, judge robe",
+    "401": "accordion, piano accordion, squeeze box",
+    "402": "acoustic guitar",
+    "403": "aircraft carrier, carrier, flattop, attack aircraft carrier",
+    "404": "airliner",
+    "405": "airship, dirigible",
+    "406": "altar",
+    "407": "ambulance",
+    "408": "amphibian, amphibious vehicle",
+    "409": "analog clock",
+    "410": "apiary, bee house",
+    "411": "apron",
+    "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin",
+    "413": "assault rifle, assault gun",
+    "414": "backpack, back pack, knapsack, packsack, rucksack, haversack",
+    "415": "bakery, bakeshop, bakehouse",
+    "416": "balance beam, beam",
+    "417": "balloon",
+    "418": "ballpoint, ballpoint pen, ballpen, Biro",
+    "419": "Band Aid",
+    "420": "banjo",
+    "421": "bannister, banister, balustrade, balusters, handrail",
+    "422": "barbell",
+    "423": "barber chair",
+    "424": "barbershop",
+    "425": "barn",
+    "426": "barometer",
+    "427": "barrel, cask",
+    "428": "barrow, garden cart, lawn cart, wheelbarrow",
+    "429": "baseball",
+    "430": "basketball",
+    "431": "bassinet",
+    "432": "bassoon",
+    "433": "bathing cap, swimming cap",
+    "434": "bath towel",
+    "435": "bathtub, bathing tub, bath, tub",
+    "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon",
+    "437": "beacon, lighthouse, beacon light, pharos",
+    "438": "beaker",
+    "439": "bearskin, busby, shako",
+    "440": "beer bottle",
+    "441": "beer glass",
+    "442": "bell cote, bell cot",
+    "443": "bib",
+    "444": "bicycle-built-for-two, tandem bicycle, tandem",
+    "445": "bikini, two-piece",
+    "446": "binder, ring-binder",
+    "447": "binoculars, field glasses, opera glasses",
+    "448": "birdhouse",
+    "449": "boathouse",
+    "450": "bobsled, bobsleigh, bob",
+    "451": "bolo tie, bolo, bola tie, bola",
+    "452": "bonnet, poke bonnet",
+    "453": "bookcase",
+    "454": "bookshop, bookstore, bookstall",
+    "455": "bottlecap",
+    "456": "bow",
+    "457": "bow tie, bow-tie, bowtie",
+    "458": "brass, memorial tablet, plaque",
+    "459": "brassiere, bra, bandeau",
+    "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty",
+    "461": "breastplate, aegis, egis",
+    "462": "broom",
+    "463": "bucket, pail",
+    "464": "buckle",
+    "465": "bulletproof vest",
+    "466": "bullet train, bullet",
+    "467": "butcher shop, meat market",
+    "468": "cab, hack, taxi, taxicab",
+    "469": "caldron, cauldron",
+    "470": "candle, taper, wax light",
+    "471": "cannon",
+    "472": "canoe",
+    "473": "can opener, tin opener",
+    "474": "cardigan",
+    "475": "car mirror",
+    "476": "carousel, carrousel, merry-go-round, roundabout, whirligig",
+    "477": "carpenters kit, tool kit",
+    "478": "carton",
+    "479": "car wheel",
+    "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM",
+    "481": "cassette",
+    "482": "cassette player",
+    "483": "castle",
+    "484": "catamaran",
+    "485": "CD player",
+    "486": "cello, violoncello",
+    "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone",
+    "488": "chain",
+    "489": "chainlink fence",
+    "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour",
+    "491": "chain saw, chainsaw",
+    "492": "chest",
+    "493": "chiffonier, commode",
+    "494": "chime, bell, gong",
+    "495": "china cabinet, china closet",
+    "496": "Christmas stocking",
+    "497": "church, church building",
+    "498": "cinema, movie theater, movie theatre, movie house, picture palace",
+    "499": "cleaver, meat cleaver, chopper",
+    "500": "cliff dwelling",
+    "501": "cloak",
+    "502": "clog, geta, patten, sabot",
+    "503": "cocktail shaker",
+    "504": "coffee mug",
+    "505": "coffeepot",
+    "506": "coil, spiral, volute, whorl, helix",
+    "507": "combination lock",
+    "508": "computer keyboard, keypad",
+    "509": "confectionery, confectionary, candy store",
+    "510": "container ship, containership, container vessel",
+    "511": "convertible",
+    "512": "corkscrew, bottle screw",
+    "513": "cornet, horn, trumpet, trump",
+    "514": "cowboy boot",
+    "515": "cowboy hat, ten-gallon hat",
+    "516": "cradle",
+    "517": "crane",
+    "518": "crash helmet",
+    "519": "crate",
+    "520": "crib, cot",
+    "521": "Crock Pot",
+    "522": "croquet ball",
+    "523": "crutch",
+    "524": "cuirass",
+    "525": "dam, dike, dyke",
+    "526": "desk",
+    "527": "desktop computer",
+    "528": "dial telephone, dial phone",
+    "529": "diaper, nappy, napkin",
+    "530": "digital clock",
+    "531": "digital watch",
+    "532": "dining table, board",
+    "533": "dishrag, dishcloth",
+    "534": "dishwasher, dish washer, dishwashing machine",
+    "535": "disk brake, disc brake",
+    "536": "dock, dockage, docking facility",
+    "537": "dogsled, dog sled, dog sleigh",
+    "538": "dome",
+    "539": "doormat, welcome mat",
+    "540": "drilling platform, offshore rig",
+    "541": "drum, membranophone, tympan",
+    "542": "drumstick",
+    "543": "dumbbell",
+    "544": "Dutch oven",
+    "545": "electric fan, blower",
+    "546": "electric guitar",
+    "547": "electric locomotive",
+    "548": "entertainment center",
+    "549": "envelope",
+    "550": "espresso maker",
+    "551": "face powder",
+    "552": "feather boa, boa",
+    "553": "file, file cabinet, filing cabinet",
+    "554": "fireboat",
+    "555": "fire engine, fire truck",
+    "556": "fire screen, fireguard",
+    "557": "flagpole, flagstaff",
+    "558": "flute, transverse flute",
+    "559": "folding chair",
+    "560": "football helmet",
+    "561": "forklift",
+    "562": "fountain",
+    "563": "fountain pen",
+    "564": "four-poster",
+    "565": "freight car",
+    "566": "French horn, horn",
+    "567": "frying pan, frypan, skillet",
+    "568": "fur coat",
+    "569": "garbage truck, dustcart",
+    "570": "gasmask, respirator, gas helmet",
+    "571": "gas pump, gasoline pump, petrol pump, island dispenser",
+    "572": "goblet",
+    "573": "go-kart",
+    "574": "golf ball",
+    "575": "golfcart, golf cart",
+    "576": "gondola",
+    "577": "gong, tam-tam",
+    "578": "gown",
+    "579": "grand piano, grand",
+    "580": "greenhouse, nursery, glasshouse",
+    "581": "grille, radiator grille",
+    "582": "grocery store, grocery, food market, market",
+    "583": "guillotine",
+    "584": "hair slide",
+    "585": "hair spray",
+    "586": "half track",
+    "587": "hammer",
+    "588": "hamper",
+    "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier",
+    "590": "hand-held computer, hand-held microcomputer",
+    "591": "handkerchief, hankie, hanky, hankey",
+    "592": "hard disc, hard disk, fixed disk",
+    "593": "harmonica, mouth organ, harp, mouth harp",
+    "594": "harp",
+    "595": "harvester, reaper",
+    "596": "hatchet",
+    "597": "holster",
+    "598": "home theater, home theatre",
+    "599": "honeycomb",
+    "600": "hook, claw",
+    "601": "hoopskirt, crinoline",
+    "602": "horizontal bar, high bar",
+    "603": "horse cart, horse-cart",
+    "604": "hourglass",
+    "605": "iPod",
+    "606": "iron, smoothing iron",
+    "607": "jack-o-lantern",
+    "608": "jean, blue jean, denim",
+    "609": "jeep, landrover",
+    "610": "jersey, T-shirt, tee shirt",
+    "611": "jigsaw puzzle",
+    "612": "jinrikisha, ricksha, rickshaw",
+    "613": "joystick",
+    "614": "kimono",
+    "615": "knee pad",
+    "616": "knot",
+    "617": "lab coat, laboratory coat",
+    "618": "ladle",
+    "619": "lampshade, lamp shade",
+    "620": "laptop, laptop computer",
+    "621": "lawn mower, mower",
+    "622": "lens cap, lens cover",
+    "623": "letter opener, paper knife, paperknife",
+    "624": "library",
+    "625": "lifeboat",
+    "626": "lighter, light, igniter, ignitor",
+    "627": "limousine, limo",
+    "628": "liner, ocean liner",
+    "629": "lipstick, lip rouge",
+    "630": "Loafer",
+    "631": "lotion",
+    "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system",
+    "633": "loupe, jewelers loupe",
+    "634": "lumbermill, sawmill",
+    "635": "magnetic compass",
+    "636": "mailbag, postbag",
+    "637": "mailbox, letter box",
+    "638": "maillot",
+    "639": "maillot, tank suit",
+    "640": "manhole cover",
+    "641": "maraca",
+    "642": "marimba, xylophone",
+    "643": "mask",
+    "644": "matchstick",
+    "645": "maypole",
+    "646": "maze, labyrinth",
+    "647": "measuring cup",
+    "648": "medicine chest, medicine cabinet",
+    "649": "megalith, megalithic structure",
+    "650": "microphone, mike",
+    "651": "microwave, microwave oven",
+    "652": "military uniform",
+    "653": "milk can",
+    "654": "minibus",
+    "655": "miniskirt, mini",
+    "656": "minivan",
+    "657": "missile",
+    "658": "mitten",
+    "659": "mixing bowl",
+    "660": "mobile home, manufactured home",
+    "661": "Model T",
+    "662": "modem",
+    "663": "monastery",
+    "664": "monitor",
+    "665": "moped",
+    "666": "mortar",
+    "667": "mortarboard",
+    "668": "mosque",
+    "669": "mosquito net",
+    "670": "motor scooter, scooter",
+    "671": "mountain bike, all-terrain bike, off-roader",
+    "672": "mountain tent",
+    "673": "mouse, computer mouse",
+    "674": "mousetrap",
+    "675": "moving van",
+    "676": "muzzle",
+    "677": "nail",
+    "678": "neck brace",
+    "679": "necklace",
+    "680": "nipple",
+    "681": "notebook, notebook computer",
+    "682": "obelisk",
+    "683": "oboe, hautboy, hautbois",
+    "684": "ocarina, sweet potato",
+    "685": "odometer, hodometer, mileometer, milometer",
+    "686": "oil filter",
+    "687": "organ, pipe organ",
+    "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO",
+    "689": "overskirt",
+    "690": "oxcart",
+    "691": "oxygen mask",
+    "692": "packet",
+    "693": "paddle, boat paddle",
+    "694": "paddlewheel, paddle wheel",
+    "695": "padlock",
+    "696": "paintbrush",
+    "697": "pajama, pyjama, pjs, jammies",
+    "698": "palace",
+    "699": "panpipe, pandean pipe, syrinx",
+    "700": "paper towel",
+    "701": "parachute, chute",
+    "702": "parallel bars, bars",
+    "703": "park bench",
+    "704": "parking meter",
+    "705": "passenger car, coach, carriage",
+    "706": "patio, terrace",
+    "707": "pay-phone, pay-station",
+    "708": "pedestal, plinth, footstall",
+    "709": "pencil box, pencil case",
+    "710": "pencil sharpener",
+    "711": "perfume, essence",
+    "712": "Petri dish",
+    "713": "photocopier",
+    "714": "pick, plectrum, plectron",
+    "715": "pickelhaube",
+    "716": "picket fence, paling",
+    "717": "pickup, pickup truck",
+    "718": "pier",
+    "719": "piggy bank, penny bank",
+    "720": "pill bottle",
+    "721": "pillow",
+    "722": "ping-pong ball",
+    "723": "pinwheel",
+    "724": "pirate, pirate ship",
+    "725": "pitcher, ewer",
+    "726": "plane, carpenters plane, woodworking plane",
+    "727": "planetarium",
+    "728": "plastic bag",
+    "729": "plate rack",
+    "730": "plow, plough",
+    "731": "plunger, plumbers helper",
+    "732": "Polaroid camera, Polaroid Land camera",
+    "733": "pole",
+    "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria",
+    "735": "poncho",
+    "736": "pool table, billiard table, snooker table",
+    "737": "pop bottle, soda bottle",
+    "738": "pot, flowerpot",
+    "739": "potters wheel",
+    "740": "power drill",
+    "741": "prayer rug, prayer mat",
+    "742": "printer",
+    "743": "prison, prison house",
+    "744": "projectile, missile",
+    "745": "projector",
+    "746": "puck, hockey puck",
+    "747": "punching bag, punch bag, punching ball, punchball",
+    "748": "purse",
+    "749": "quill, quill pen",
+    "750": "quilt, comforter, comfort, puff",
+    "751": "racer, race car, racing car",
+    "752": "racket, racquet",
+    "753": "radiator",
+    "754": "radio, wireless",
+    "755": "radio telescope, radio reflector",
+    "756": "rain barrel",
+    "757": "recreational vehicle, RV, R.V.",
+    "758": "reel",
+    "759": "reflex camera",
+    "760": "refrigerator, icebox",
+    "761": "remote control, remote",
+    "762": "restaurant, eating house, eating place, eatery",
+    "763": "revolver, six-gun, six-shooter",
+    "764": "rifle",
+    "765": "rocking chair, rocker",
+    "766": "rotisserie",
+    "767": "rubber eraser, rubber, pencil eraser",
+    "768": "rugby ball",
+    "769": "rule, ruler",
+    "770": "running shoe",
+    "771": "safe",
+    "772": "safety pin",
+    "773": "saltshaker, salt shaker",
+    "774": "sandal",
+    "775": "sarong",
+    "776": "sax, saxophone",
+    "777": "scabbard",
+    "778": "scale, weighing machine",
+    "779": "school bus",
+    "780": "schooner",
+    "781": "scoreboard",
+    "782": "screen, CRT screen",
+    "783": "screw",
+    "784": "screwdriver",
+    "785": "seat belt, seatbelt",
+    "786": "sewing machine",
+    "787": "shield, buckler",
+    "788": "shoe shop, shoe-shop, shoe store",
+    "789": "shoji",
+    "790": "shopping basket",
+    "791": "shopping cart",
+    "792": "shovel",
+    "793": "shower cap",
+    "794": "shower curtain",
+    "795": "ski",
+    "796": "ski mask",
+    "797": "sleeping bag",
+    "798": "slide rule, slipstick",
+    "799": "sliding door",
+    "800": "slot, one-armed bandit",
+    "801": "snorkel",
+    "802": "snowmobile",
+    "803": "snowplow, snowplough",
+    "804": "soap dispenser",
+    "805": "soccer ball",
+    "806": "sock",
+    "807": "solar dish, solar collector, solar furnace",
+    "808": "sombrero",
+    "809": "soup bowl",
+    "810": "space bar",
+    "811": "space heater",
+    "812": "space shuttle",
+    "813": "spatula",
+    "814": "speedboat",
+    "815": "spider web, spiders web",
+    "816": "spindle",
+    "817": "sports car, sport car",
+    "818": "spotlight, spot",
+    "819": "stage",
+    "820": "steam locomotive",
+    "821": "steel arch bridge",
+    "822": "steel drum",
+    "823": "stethoscope",
+    "824": "stole",
+    "825": "stone wall",
+    "826": "stopwatch, stop watch",
+    "827": "stove",
+    "828": "strainer",
+    "829": "streetcar, tram, tramcar, trolley, trolley car",
+    "830": "stretcher",
+    "831": "studio couch, day bed",
+    "832": "stupa, tope",
+    "833": "submarine, pigboat, sub, U-boat",
+    "834": "suit, suit of clothes",
+    "835": "sundial",
+    "836": "sunglass",
+    "837": "sunglasses, dark glasses, shades",
+    "838": "sunscreen, sunblock, sun blocker",
+    "839": "suspension bridge",
+    "840": "swab, swob, mop",
+    "841": "sweatshirt",
+    "842": "swimming trunks, bathing trunks",
+    "843": "swing",
+    "844": "switch, electric switch, electrical switch",
+    "845": "syringe",
+    "846": "table lamp",
+    "847": "tank, army tank, armored combat vehicle, armoured combat vehicle",
+    "848": "tape player",
+    "849": "teapot",
+    "850": "teddy, teddy bear",
+    "851": "television, television system",
+    "852": "tennis ball",
+    "853": "thatch, thatched roof",
+    "854": "theater curtain, theatre curtain",
+    "855": "thimble",
+    "856": "thresher, thrasher, threshing machine",
+    "857": "throne",
+    "858": "tile roof",
+    "859": "toaster",
+    "860": "tobacco shop, tobacconist shop, tobacconist",
+    "861": "toilet seat",
+    "862": "torch",
+    "863": "totem pole",
+    "864": "tow truck, tow car, wrecker",
+    "865": "toyshop",
+    "866": "tractor",
+    "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi",
+    "868": "tray",
+    "869": "trench coat",
+    "870": "tricycle, trike, velocipede",
+    "871": "trimaran",
+    "872": "tripod",
+    "873": "triumphal arch",
+    "874": "trolleybus, trolley coach, trackless trolley",
+    "875": "trombone",
+    "876": "tub, vat",
+    "877": "turnstile",
+    "878": "typewriter keyboard",
+    "879": "umbrella",
+    "880": "unicycle, monocycle",
+    "881": "upright, upright piano",
+    "882": "vacuum, vacuum cleaner",
+    "883": "vase",
+    "884": "vault",
+    "885": "velvet",
+    "886": "vending machine",
+    "887": "vestment",
+    "888": "viaduct",
+    "889": "violin, fiddle",
+    "890": "volleyball",
+    "891": "waffle iron",
+    "892": "wall clock",
+    "893": "wallet, billfold, notecase, pocketbook",
+    "894": "wardrobe, closet, press",
+    "895": "warplane, military plane",
+    "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin",
+    "897": "washer, automatic washer, washing machine",
+    "898": "water bottle",
+    "899": "water jug",
+    "900": "water tower",
+    "901": "whiskey jug",
+    "902": "whistle",
+    "903": "wig",
+    "904": "window screen",
+    "905": "window shade",
+    "906": "Windsor tie",
+    "907": "wine bottle",
+    "908": "wing",
+    "909": "wok",
+    "910": "wooden spoon",
+    "911": "wool, woolen, woollen",
+    "912": "worm fence, snake fence, snake-rail fence, Virginia fence",
+    "913": "wreck",
+    "914": "yawl",
+    "915": "yurt",
+    "916": "web site, website, internet site, site",
+    "917": "comic book",
+    "918": "crossword puzzle, crossword",
+    "919": "street sign",
+    "920": "traffic light, traffic signal, stoplight",
+    "921": "book jacket, dust cover, dust jacket, dust wrapper",
+    "922": "menu",
+    "923": "plate",
+    "924": "guacamole",
+    "925": "consomme",
+    "926": "hot pot, hotpot",
+    "927": "trifle",
+    "928": "ice cream, icecream",
+    "929": "ice lolly, lolly, lollipop, popsicle",
+    "930": "French loaf",
+    "931": "bagel, beigel",
+    "932": "pretzel",
+    "933": "cheeseburger",
+    "934": "hotdog, hot dog, red hot",
+    "935": "mashed potato",
+    "936": "head cabbage",
+    "937": "broccoli",
+    "938": "cauliflower",
+    "939": "zucchini, courgette",
+    "940": "spaghetti squash",
+    "941": "acorn squash",
+    "942": "butternut squash",
+    "943": "cucumber, cuke",
+    "944": "artichoke, globe artichoke",
+    "945": "bell pepper",
+    "946": "cardoon",
+    "947": "mushroom",
+    "948": "Granny Smith",
+    "949": "strawberry",
+    "950": "orange",
+    "951": "lemon",
+    "952": "fig",
+    "953": "pineapple, ananas",
+    "954": "banana",
+    "955": "jackfruit, jak, jack",
+    "956": "custard apple",
+    "957": "pomegranate",
+    "958": "hay",
+    "959": "carbonara",
+    "960": "chocolate sauce, chocolate syrup",
+    "961": "dough",
+    "962": "meat loaf, meatloaf",
+    "963": "pizza, pizza pie",
+    "964": "potpie",
+    "965": "burrito",
+    "966": "red wine",
+    "967": "espresso",
+    "968": "cup",
+    "969": "eggnog",
+    "970": "alp",
+    "971": "bubble",
+    "972": "cliff, drop, drop-off",
+    "973": "coral reef",
+    "974": "geyser",
+    "975": "lakeside, lakeshore",
+    "976": "promontory, headland, head, foreland",
+    "977": "sandbar, sand bar",
+    "978": "seashore, coast, seacoast, sea-coast",
+    "979": "valley, vale",
+    "980": "volcano",
+    "981": "ballplayer, baseball player",
+    "982": "groom, bridegroom",
+    "983": "scuba diver",
+    "984": "rapeseed",
+    "985": "daisy",
+    "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",
+    "987": "corn",
+    "988": "acorn",
+    "989": "hip, rose hip, rosehip",
+    "990": "buckeye, horse chestnut, conker",
+    "991": "coral fungus",
+    "992": "agaric",
+    "993": "gyromitra",
+    "994": "stinkhorn, carrion fungus",
+    "995": "earthstar",
+    "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa",
+    "997": "bolete",
+    "998": "ear, spike, capitulum",
+    "999": "toilet tissue, toilet paper, bathroom tissue"
+  }
+}

W-Flow-B/pipeline.py

@@ -0,0 +1,754 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.util
+import inspect
+import json
+import sys
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import deprecate, is_torch_xla_available, logging, replace_example_docstring
+from PIL import Image
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from diffusers import DiffusionPipeline
+
+        >>> model_dir = Path("W-Flow-diffusers/W-Flow-B").resolve()
+        >>> pipe = DiffusionPipeline.from_pretrained(
+        ...     str(model_dir),
+        ...     local_files_only=True,
+        ...     custom_pipeline=str(model_dir / "pipeline.py"),
+        ...     trust_remote_code=True,
+        ...     torch_dtype=torch.bfloat16,
+        ... ).to("cuda")
+
+        >>> image = pipe(
+        ...     class_labels="golden retriever",
+        ...     num_inference_steps=1,
+        ...     guidance_scale=1.19,
+        ...     generator=torch.Generator(device="cuda").manual_seed(1),
+        ... ).images[0]
+        >>> image.save("demo.png")
+
+        >>> # Hugging Face Hub style model id: UserID/RepoID
+        >>> # pipe = DiffusionPipeline.from_pretrained(
+        >>> #     "BiliSakura/W-Flow-diffusers",
+        >>> #     subfolder="W-Flow-B",
+        >>> #     custom_pipeline="pipeline.py",
+        >>> #     trust_remote_code=True,
+        >>> #     torch_dtype=torch.bfloat16,
+        >>> # )
+        ```
+"""
+
+
+class WFlowPipeline(DiffusionPipeline):
+    r"""
+    Class-conditional one-step pipeline for W-Flow.
+
+    W-Flow maps SD-VAE latent noise to a clean 32×32 latent in one forward pass, then decodes it
+    with `stabilityai/sd-vae-ft-mse` to a 256×256 image. Guidance is embedded in the generator
+    (the official `cfg_scale`, which is `w + 1` in the paper). There is no unconditional class
+    and no second forward.
+
+    The pipeline inherits from [`DiffusionPipeline`]. Check the superclass documentation for the
+    generic methods (download, save, device placement, CPU offload).
+
+    Args:
+        transformer ([`ModelMixin`]):
+            [`WFlowTransformer2DModel`] that predicts the clean latent.
+        vae ([`AutoencoderKL`]):
+            SD-VAE decoder. Released checkpoints use `stabilityai/sd-vae-ft-mse` and scaling 0.18215.
+        scheduler ([`WFlowScheduler`]):
+            One-step scheduler. Loaded from `scheduler/`.
+        id2label (`dict[int, str]`, *optional*):
+            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
+            Loaded from `model_index.json` when omitted.
+
+        The default guidance scale lives on `transformer.config.guidance_scale`.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+    _callback_tensor_inputs = ["latents"]
+
+    def __init__(
+        self,
+        transformer: ModelMixin,
+        vae: AutoencoderKL,
+        scheduler: SchedulerMixin,
+        id2label: dict[int, str] | dict[str, str] | None = None,
+    ):
+        super().__init__()
+        if scheduler is None:
+            scheduler = self._default_scheduler()
+        self._reject_karras_scheduler(scheduler)
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+        normalized_labels = self._normalize_id2label(id2label)
+        self.register_to_config(id2label={str(key): value for key, value in normalized_labels.items()})
+        self._id2label = normalized_labels
+        self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = bool(self._id2label)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    def _sampling_default(self, name: str, fallback: float) -> float:
+        r"""
+        Read a sampling default from the transformer config.
+
+        Args:
+            name (`str`):
+                Key on `transformer.config`, such as `guidance_scale`.
+            fallback (`float`):
+                Value used when the transformer config does not define `name`.
+
+        Returns:
+            `float`: The configured value, or `fallback`.
+        """
+        value = getattr(self.transformer.config, name, None)
+        if value is None:
+            return fallback
+        return value
+
+    @staticmethod
+    def _reject_karras_scheduler(scheduler: SchedulerMixin) -> None:
+        r"""
+        Reject multi-step Karras schedulers.
+
+        W-Flow predicts the clean latent in one evaluation and has no diffusion timestep, so a
+        Karras noise schedule cannot be integrated.
+
+        Args:
+            scheduler ([`SchedulerMixin`]):
+                Scheduler attached to the pipeline.
+        """
+        scheduler_name = scheduler.__class__.__name__
+        if scheduler_name in KarrasDiffusionSchedulers.__members__:
+            raise TypeError(
+                "W-Flow is a one-step generator and does not integrate a Karras noise schedule. "
+                f"Received {scheduler_name}. Use WFlowScheduler."
+            )
+
+    @staticmethod
+    def _default_scheduler() -> SchedulerMixin:
+        r"""
+        Build [`WFlowScheduler`] from the variant directory.
+
+        Returns:
+            [`WFlowScheduler`]: One-step sample scheduler.
+        """
+        module_path = Path(__file__).resolve().parent / "scheduler" / "scheduling_wflow.py"
+        spec = importlib.util.spec_from_file_location("scheduling_wflow", module_path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Could not load the W-Flow scheduler from {module_path}.")
+        module = importlib.util.module_from_spec(spec)
+        sys.modules["scheduling_wflow"] = module
+        spec.loader.exec_module(module)
+        return module.WFlowScheduler()
+
+    def _ensure_labels_loaded(self) -> None:
+        r"""
+        Fill `id2label` from `model_index.json` when the constructor did not receive it.
+        """
+        if self._labels_loaded_from_model_index:
+            return
+        loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None))
+        if loaded:
+            self._id2label = loaded
+            self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = True
+
+    @staticmethod
+    def _normalize_id2label(id2label: dict[int, str] | dict[str, str] | None) -> dict[int, str]:
+        r"""
+        Cast label-map keys to integers.
+
+        Args:
+            id2label (`dict`, *optional*):
+                Class id to English label. Keys may be strings, which is how JSON stores them.
+
+        Returns:
+            `dict[int, str]`: Integer-keyed map. Empty when `id2label` is missing.
+        """
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _read_id2label_from_model_index(variant_path: str | None) -> dict[int, str]:
+        r"""
+        Read `id2label` from a local variant directory.
+
+        Args:
+            variant_path (`str`, *optional*):
+                Directory that contains `model_index.json`.
+
+        Returns:
+            `dict[int, str]`: Loaded map, or an empty dict when the file or key is absent.
+        """
+        if not variant_path:
+            return {}
+        model_index_path = Path(variant_path).resolve() / "model_index.json"
+        if not model_index_path.exists():
+            return {}
+        raw = json.loads(model_index_path.read_text(encoding="utf-8"))
+        id2label = raw.get("id2label")
+        if not isinstance(id2label, dict):
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _build_label2id(id2label: dict[int, str]) -> dict[str, int]:
+        r"""
+        Invert `id2label`, splitting comma-separated synonyms.
+
+        Args:
+            id2label (`dict[int, str]`):
+                Class id to one or more English names.
+
+        Returns:
+            `dict[str, int]`: Synonym to class id, sorted by synonym.
+        """
+        label2id: dict[str, int] = {}
+        for class_id, value in id2label.items():
+            for synonym in value.split(","):
+                synonym = synonym.strip()
+                if synonym:
+                    label2id[synonym] = int(class_id)
+        return dict(sorted(label2id.items()))
+
+    @property
+    def interrupt(self) -> bool:
+        r"""
+        Whether the current call was asked to stop.
+
+        Returns:
+            `bool`: The `_interrupt` flag. `__call__` sets it to `False` at the start. Assign
+            `_interrupt = True` during a call to skip the remaining forward.
+        """
+        return self._interrupt
+
+    @property
+    def id2label(self) -> dict[int, str]:
+        r"""
+        ImageNet id to English label map.
+
+        Returns:
+            `dict[int, str]`: The map loaded from the constructor or `model_index.json`.
+        """
+        self._ensure_labels_loaded()
+        return self._id2label
+
+    def get_label_ids(self, label: str | list[str]) -> list[int]:
+        r"""
+        Map ImageNet English labels to class ids.
+
+        Args:
+            label (`str` or `list[str]`):
+                One label or a list of labels. Each string must match a synonym in `id2label`
+                after stripping whitespace. Synonyms are the comma-separated pieces of each value.
+
+        Returns:
+            `list[int]`: Class ids in the same order as `label`.
+        """
+        self._ensure_labels_loaded()
+        label2id = self.labels
+        if not label2id:
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
+        if isinstance(label, str):
+            label = [label]
+        missing = [item for item in label if item not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
+        return [label2id[item] for item in label]
+
+    def _normalize_class_labels(self, class_labels: int | str | list[int | str]) -> list[int]:
+        r"""
+        Turn a class-label argument into a list of integer ids.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                One id, one English label, or a list of either. A list must not mix types.
+
+        Returns:
+            `list[int]`: One id per requested image.
+        """
+        if isinstance(class_labels, int):
+            return [class_labels]
+        if isinstance(class_labels, str):
+            return self.get_label_ids(class_labels)
+        if len(class_labels) == 0:
+            raise ValueError("`class_labels` must contain at least one class id or label.")
+        if isinstance(class_labels[0], str):
+            if not all(isinstance(item, str) for item in class_labels):
+                raise TypeError("`class_labels` must be all strings or all integers.")
+            return self.get_label_ids(class_labels)
+        if not all(isinstance(item, int) for item in class_labels):
+            raise TypeError("`class_labels` must be all strings or all integers.")
+        return list(class_labels)
+
+    def prepare_extra_step_kwargs(
+        self,
+        generator: torch.Generator | list[torch.Generator] | None,
+        eta: float,
+    ) -> dict[str, Any]:
+        r"""
+        Extra arguments forwarded to `scheduler.step`, depending on what that method accepts.
+
+        Args:
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Generator passed through when the scheduler step samples noise.
+            eta (`float`):
+                DDIM eta in `[0, 1]`. Ignored by [`WFlowScheduler`], whose `step` has no `eta` argument.
+
+        Returns:
+            `dict`: Keyword arguments for `scheduler.step`.
+        """
+        extra_step_kwargs: dict[str, Any] = {}
+        step_params = set(inspect.signature(self.scheduler.step).parameters.keys())
+        if "eta" in step_params:
+            extra_step_kwargs["eta"] = eta
+        if "generator" in step_params:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        class_labels: int | str | list[int | str],
+        height: int,
+        width: int,
+        num_inference_steps: int,
+        guidance_scale: float,
+        noise_scale: float,
+        output_type: str,
+        callback_steps: int | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+    ) -> None:
+        r"""
+        Reject arguments that the generator cannot run.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                Class ids or English labels. A list must be non-empty.
+            height (`int`):
+                Output height in pixels. Must be the pretrained resolution, 256.
+            width (`int`):
+                Output width in pixels. Must be the pretrained resolution, 256.
+            num_inference_steps (`int`):
+                Must be 1. W-Flow evaluates the network once.
+            guidance_scale (`float`):
+                Embedded guidance (`w + 1` in the paper). Must be finite and positive.
+            noise_scale (`float`):
+                Scale of the reference Gaussian, the official `temp`. Must be positive.
+            output_type (`str`):
+                One of `"pil"`, `"np"`, `"pt"`, `"latent"`.
+            callback_steps (`int`, *optional*):
+                Deprecated callback period. When set, must be a positive integer.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names the step callback may read. Each name must be in `_callback_tensor_inputs`.
+        """
+        if isinstance(class_labels, list) and len(class_labels) == 0:
+            raise ValueError("`class_labels` must contain at least one class id or label.")
+        native = int(self.transformer.config.sample_size) * int(self.vae_scale_factor)
+        if height != native or width != native:
+            raise ValueError(
+                f"`height` and `width` must be the pretrained resolution {native}. "
+                f"The patch embedding is fixed to a {self.transformer.config.sample_size} latent. Got {(height, width)}."
+            )
+        if num_inference_steps != 1:
+            raise ValueError(
+                "W-Flow generates in one forward pass. "
+                f"`num_inference_steps` must be 1, got {num_inference_steps}."
+            )
+        if not torch.isfinite(torch.tensor(float(guidance_scale))) or float(guidance_scale) <= 0:
+            raise ValueError(f"`guidance_scale` must be finite and positive, got {guidance_scale}.")
+        if noise_scale <= 0:
+            raise ValueError(f"`noise_scale` must be positive, got {noise_scale}.")
+        if output_type not in {"pil", "np", "pt", "latent"}:
+            raise ValueError(f"Unsupported `output_type` {output_type!r}. Choose from 'pil', 'np', 'pt', 'latent'.")
+        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps < 1):
+            raise ValueError(f"`callback_steps` must be a positive integer, got {callback_steps}.")
+        if callback_on_step_end_tensor_inputs is not None:
+            unknown = [name for name in callback_on_step_end_tensor_inputs if name not in self._callback_tensor_inputs]
+            if unknown:
+                raise ValueError(
+                    f"`callback_on_step_end_tensor_inputs` contains {unknown}, which are not in "
+                    f"{self._callback_tensor_inputs}."
+                )
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels: int,
+        height: int,
+        width: int,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None,
+        noise_labels: torch.Tensor | None,
+        noise_scale: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Draw reference noise in the official generator order.
+
+        The official sampler draws `(batch, height, width, channels)` first, then integer noise
+        codes, from the same generator. This method follows that order so a seeded generator
+        matches `DitGen.forward`.
+
+        Args:
+            batch_size (`int`):
+                Number of images.
+            num_channels (`int`):
+                Latent channels. Released checkpoints use 4.
+            height (`int`):
+                Latent height. Released checkpoints use 32.
+            width (`int`):
+                Latent width. Released checkpoints use 32.
+            device (`torch.device`):
+                Device of the returned tensors.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG. A list must have one generator per image. Each generator is consumed in the
+                order Gaussian noise, then noise codes.
+            latents (`torch.Tensor`, *optional*):
+                Reference noise `(batch, channels, height, width)`. When set, it is not rescaled.
+            noise_labels (`torch.Tensor`, *optional*):
+                Integer codes `(batch, noise_coords)`. When omitted, they are drawn here.
+            noise_scale (`float`):
+                Multiplier on the Gaussian, the official `temp`. Ignored when `latents` is passed.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`:
+                Reference noise in BCHW float32, and noise codes `(batch, noise_coords)`.
+        """
+        noise_classes = int(self.transformer.config.noise_classes)
+        noise_coords = int(self.transformer.config.noise_coords)
+        if latents is None:
+            if isinstance(generator, list):
+                if len(generator) != batch_size:
+                    raise ValueError(f"Got {len(generator)} generators for a batch of {batch_size}.")
+                pieces = [
+                    torch.randn((1, height, width, num_channels), generator=one, device=device) for one in generator
+                ]
+                drawn = torch.cat(pieces, dim=0)
+            else:
+                drawn = torch.randn(
+                    (batch_size, height, width, num_channels),
+                    generator=generator,
+                    device=device,
+                )
+            latents = drawn.permute(0, 3, 1, 2).contiguous() * float(noise_scale)
+        elif latents.shape != (batch_size, num_channels, height, width):
+            raise ValueError(
+                f"`latents` must have shape {(batch_size, num_channels, height, width)}, got {tuple(latents.shape)}."
+            )
+        else:
+            latents = latents.to(device=device)
+
+        if noise_labels is None:
+            if isinstance(generator, list) and latents is not None:
+                # When latents were supplied, generators have not been consumed yet only if the
+                # caller passed both. Drawing labels still uses those generators.
+                codes = [
+                    torch.randint(0, noise_classes, (1, noise_coords), generator=one, device=device)
+                    for one in generator
+                ]
+                noise_labels = torch.cat(codes, dim=0)
+            else:
+                noise_labels = torch.randint(
+                    0,
+                    max(1, noise_classes),
+                    (batch_size, max(1, noise_coords)),
+                    generator=None if isinstance(generator, list) else generator,
+                    device=device,
+                )
+        noise_labels = noise_labels.to(device=device, dtype=torch.long)
+        if noise_labels.shape != (batch_size, noise_coords):
+            raise ValueError(
+                f"`noise_labels` must have shape {(batch_size, noise_coords)}, got {tuple(noise_labels.shape)}."
+            )
+        return latents.float(), noise_labels
+
+    def _prepare_scheduler(self, num_inference_steps: int, device: torch.device) -> torch.Tensor:
+        r"""
+        Install the one-step schedule and return its timesteps.
+
+        Args:
+            num_inference_steps (`int`):
+                Must be 1.
+            device (`torch.device`):
+                Device of the timestep tensor.
+
+        Returns:
+            `torch.Tensor`: The single timestep.
+        """
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        return self.scheduler.timesteps
+
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        r"""
+        Decode scaled SD-VAE latents to pixels in `[-1, 1]`.
+
+        Args:
+            latents (`torch.Tensor`):
+                Clean latents `(batch, 4, 32, 32)` in the training scale (`* 0.18215`).
+
+        Returns:
+            `torch.Tensor`: Images `(batch, 3, height, width)` in `[-1, 1]`.
+        """
+        scaling = float(getattr(self.vae.config, "scaling_factor", 0.18215))
+        needs_upcast = bool(getattr(self.vae.config, "force_upcast", False)) and self.vae.dtype != torch.float32
+        original_dtype = self.vae.dtype
+        if needs_upcast:
+            self.vae.to(dtype=torch.float32)
+        decoded = self.vae.decode((latents / scaling).to(dtype=self.vae.dtype), return_dict=False)[0]
+        if needs_upcast:
+            self.vae.to(dtype=original_dtype)
+        return decoded
+
+    def _images_from_pixels(self, pixels: torch.Tensor, output_type: str) -> list[Image.Image] | torch.Tensor | list:
+        r"""
+        Convert decoded pixels to the requested output type.
+
+        Args:
+            pixels (`torch.Tensor`):
+                Images `(batch, 3, height, width)` in `[-1, 1]`.
+            output_type (`str`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"`.
+
+        Returns:
+            Images in the requested type.
+        """
+        images = ((pixels.float() + 1) / 2).clamp(0, 1)
+        if output_type == "pt":
+            return images
+        images = (images * 255).round().to(torch.uint8)
+        arrays = images.permute(0, 2, 3, 1).cpu().numpy()
+        if output_type == "np":
+            return arrays
+        return [Image.fromarray(image) for image in arrays]
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        class_labels: int | str | list[int | str],
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        guidance_scale: float | None = None,
+        noise_scale: float = 1.0,
+        eta: float = 0.0,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        noise_labels: torch.Tensor | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Callable[[int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+        **kwargs,
+    ) -> ImagePipelineOutput | tuple:
+        r"""
+        Generate images from ImageNet class labels.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                One ImageNet class id, one English label, or a list of either. English labels are
+                resolved with [`get_label_ids`]. The list length is the batch size.
+            height (`int`, *optional*):
+                Image height in pixels. Defaults to 256 (`sample_size * vae_scale_factor`). Other
+                sizes are rejected because the patch embedding is fixed.
+            width (`int`, *optional*):
+                Image width in pixels. Defaults to 256.
+            num_inference_steps (`int`, *optional*):
+                Must be 1. Defaults to `transformer.config.default_num_inference_steps`.
+            guidance_scale (`float`, *optional*):
+                Official CFG value, equal to `w + 1` in the paper. Defaults to
+                `transformer.config.guidance_scale` (1.19 / 1.14 / 1.09 / 1.5 by variant). This
+                scale is embedded; the pipeline does not run an unconditional branch.
+            noise_scale (`float`, defaults to 1.0):
+                Multiplier on the reference Gaussian. The official sampler calls this `temp`.
+                Ignored when `latents` is passed.
+            eta (`float`, defaults to 0.0):
+                DDIM eta. Ignored by [`WFlowScheduler`].
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG for the reference noise and the discrete noise codes. One generator is consumed
+                in that order, matching the official sampler.
+            latents (`torch.Tensor`, *optional*):
+                Reference noise `(batch, 4, 32, 32)`. Not multiplied by `noise_scale`.
+            noise_labels (`torch.Tensor`, *optional*):
+                Discrete noise codes `(batch, noise_coords)`. Drawn from `generator` when omitted.
+            output_type (`str`, defaults to `"pil"`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"` (scaled SD-VAE latent).
+            return_dict (`bool`, defaults to `True`):
+                Return [`ImagePipelineOutput`] when `True`, otherwise a one-tuple of images.
+            callback_on_step_end (`Callable` or `PipelineCallback`, *optional*):
+                Called as `callback_on_step_end(pipeline, step, timestep, callback_kwargs)` after the
+                scheduler step. Return a dict to replace tensors listed in
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names passed to the step callback. Defaults to `["latents"]`.
+
+        Examples:
+
+        Returns:
+            [`ImagePipelineOutput`] or `tuple`:
+                When `return_dict` is `True`, [`ImagePipelineOutput`] with the images. Otherwise a tuple
+                whose first element is the images.
+        """
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        if kwargs:
+            raise TypeError(f"Unexpected arguments: {sorted(kwargs)}.")
+        if callback_on_step_end_tensor_inputs is None:
+            callback_on_step_end_tensor_inputs = ["latents"]
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        image_size = int(self.transformer.config.sample_size) * int(self.vae_scale_factor)
+        height = image_size if height is None else int(height)
+        width = image_size if width is None else int(width)
+        if num_inference_steps is None:
+            num_inference_steps = int(self._sampling_default("default_num_inference_steps", 1))
+        if guidance_scale is None:
+            guidance_scale = float(self._sampling_default("guidance_scale", 1.0))
+
+        # 1. Check inputs.
+        self.check_inputs(
+            class_labels,
+            height,
+            width,
+            num_inference_steps,
+            guidance_scale,
+            noise_scale,
+            output_type,
+            callback_steps,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = float(guidance_scale)
+        self._interrupt = False
+
+        # 2. Define call parameters.
+        class_label_ids = self._normalize_class_labels(class_labels)
+        batch_size = len(class_label_ids)
+        device = self._execution_device
+        num_classes = int(self.transformer.config.num_classes)
+        if any(class_id < 0 or class_id >= num_classes for class_id in class_label_ids):
+            raise ValueError(f"Class ids must lie in [0, {num_classes}), got {class_label_ids}.")
+
+        # 3. Encode input condition (class ids; guidance is embedded inside the transformer).
+        class_labels_t = torch.tensor(class_label_ids, device=device, dtype=torch.long)
+
+        # 4. Prepare timesteps.
+        scheduled_timesteps = self._prepare_scheduler(num_inference_steps, device)
+
+        # 5. Prepare latent variables. Noise codes are drawn here so one generator matches DitGen.
+        latent_height = height // int(self.vae_scale_factor)
+        latent_width = width // int(self.vae_scale_factor)
+        latents, noise_labels = self.prepare_latents(
+            batch_size,
+            int(self.transformer.config.in_channels),
+            latent_height,
+            latent_width,
+            device,
+            generator,
+            latents,
+            noise_labels,
+            noise_scale,
+        )
+
+        # 6. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_warmup_steps = len(scheduled_timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(scheduled_timesteps)
+
+        # 7. Denoising loop.
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, timestep in enumerate(scheduled_timesteps):
+                if self.interrupt:
+                    continue
+
+                latent_model_input = self.scheduler.scale_model_input(latents, timestep)
+                model_output = self.transformer(
+                    latent_model_input,
+                    class_labels=class_labels_t,
+                    guidance_scale=self._guidance_scale,
+                    noise_labels=noise_labels,
+                    return_dict=False,
+                )[0]
+                latents = self.scheduler.step(
+                    model_output,
+                    timestep,
+                    latents,
+                    **extra_step_kwargs,
+                    return_dict=False,
+                )[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {key: locals()[key] for key in callback_on_step_end_tensor_inputs}
+                    callback_outputs = callback_on_step_end(self, i, timestep, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+
+                if i == len(scheduled_timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, timestep, latents)
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            images = latents
+        else:
+            images = self._images_from_pixels(self._decode_latents(latents), output_type)
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

W-Flow-B/scheduler/scheduler_config.json

@@ -0,0 +1,6 @@
+{
+  "_class_name": "WFlowScheduler",
+  "_diffusers_version": "0.39.0",
+  "num_train_timesteps": 1000,
+  "prediction_type": "sample"
+}

W-Flow-B/scheduler/scheduling_wflow.py

@@ -0,0 +1,152 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class WFlowSchedulerOutput(BaseOutput):
+    r"""
+    Output of [`WFlowScheduler.step`].
+
+    Args:
+        prev_sample (`torch.Tensor`):
+            The model prediction. W-Flow returns the clean latent in one evaluation.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class WFlowScheduler(SchedulerMixin, ConfigMixin):
+    r"""
+    One-step scheduler for W-Flow.
+
+    The transformer predicts the clean latent directly. `set_timesteps` accepts a single evaluation,
+    and `step` returns that prediction. `scale_model_input` leaves the reference noise unchanged.
+    A Karras noise schedule does not apply: the network has no diffusion timestep.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            Stored for config compatibility. W-Flow does not index it.
+        prediction_type (`str`, defaults to `"sample"`):
+            The model output is the clean sample.
+    """
+
+    order = 1
+    _compatibles: list[str] = []
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        prediction_type: str = "sample",
+    ):
+        if num_train_timesteps < 1:
+            raise ValueError(f"num_train_timesteps must be >= 1, got {num_train_timesteps}.")
+        if prediction_type != "sample":
+            raise ValueError(f"WFlowScheduler prediction_type must be 'sample', got {prediction_type!r}.")
+        self.timesteps: torch.Tensor | None = None
+        self.num_inference_steps: int | None = None
+        self._step_index: int | None = None
+
+    @property
+    def init_noise_sigma(self) -> float:
+        r"""
+        Scale of the standard normal used to start sampling.
+
+        Returns:
+            `float`: `1.0`. The pipeline multiplies this by `noise_scale`.
+        """
+        return 1.0
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
+        r"""
+        Return the sample unchanged.
+
+        Args:
+            sample (`torch.Tensor`):
+                Reference noise `(batch, channels, height, width)`.
+            timestep (`float` or `torch.Tensor`):
+                Unused. W-Flow has no diffusion timestep.
+
+        Returns:
+            `torch.Tensor`: `sample`.
+        """
+        del timestep
+        return sample
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: str | torch.device | None = None,
+    ) -> None:
+        r"""
+        Install the single evaluation step.
+
+        Args:
+            num_inference_steps (`int`):
+                Must be 1. W-Flow generates the latent in one forward pass.
+            device (`str` or `torch.device`, *optional*):
+                Device of `self.timesteps`.
+        """
+        if num_inference_steps != 1:
+            raise ValueError(
+                "W-Flow generates in one forward pass. "
+                f"num_inference_steps must be 1, got {num_inference_steps}."
+            )
+        self.num_inference_steps = 1
+        self.timesteps = torch.tensor([0.0], device=device, dtype=torch.float32)
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: float | torch.Tensor,
+        sample: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        return_dict: bool = True,
+    ) -> WFlowSchedulerOutput | tuple[torch.Tensor]:
+        r"""
+        Replace the reference noise with the predicted clean latent.
+
+        Args:
+            model_output (`torch.Tensor`):
+                Clean-latent prediction `(batch, channels, height, width)`.
+            timestep (`float` or `torch.Tensor`):
+                The single scheduled timestep. Accepted for the Diffusers step signature.
+            sample (`torch.Tensor`):
+                Current reference noise. Same shape as `model_output`.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Accepted so pipelines can pass it through. W-Flow does not draw noise in `step`.
+            return_dict (`bool`, defaults to `True`):
+                Return [`WFlowSchedulerOutput`] when `True`, otherwise a one-tuple.
+
+        Returns:
+            [`WFlowSchedulerOutput`] or `tuple`:
+                `prev_sample` is `model_output`.
+        """
+        del timestep, sample, generator
+        if self.timesteps is None:
+            raise ValueError("Call set_timesteps before step.")
+        if self._step_index is None or self._step_index > 0:
+            raise ValueError("WFlowScheduler has already consumed its single step. Call set_timesteps again.")
+        self._step_index = 1
+        if not return_dict:
+            return (model_output,)
+        return WFlowSchedulerOutput(prev_sample=model_output)

W-Flow-B/transformer/config.json

@@ -0,0 +1,26 @@
+{
+  "_class_name": "WFlowTransformer2DModel",
+  "_diffusers_version": "0.39.0",
+  "attn_fp32": false,
+  "cond_dim": 768,
+  "default_num_inference_steps": 1,
+  "depth": 12,
+  "guidance_scale": 1.19,
+  "hidden_size": 768,
+  "in_channels": 4,
+  "mlp_ratio": 4.0,
+  "model_type": "W-Flow-B",
+  "n_cls_tokens": 16,
+  "noise_classes": 64,
+  "noise_coords": 32,
+  "num_attention_heads": 12,
+  "num_classes": 1000,
+  "out_channels": 4,
+  "patch_size": 2,
+  "sample_size": 32,
+  "use_bf16": true,
+  "use_qknorm": true,
+  "use_rmsnorm": true,
+  "use_rope": true,
+  "use_swiglu": true
+}

W-Flow-B/transformer/diffusion_pytorch_model.safetensors

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

W-Flow-B/transformer/transformer_wflow.py

@@ -0,0 +1,720 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class WFlowTransformer2DModelOutput(BaseOutput):
+    r"""
+    Output of [`WFlowTransformer2DModel`].
+
+    Args:
+        sample (`torch.Tensor`):
+            Predicted SD-VAE latent of shape `(batch, channels, height, width)`. Values are in the
+            scaled latent space used at training time (`encoder_mean * 0.18215`).
+    """
+
+    sample: torch.Tensor
+
+
+def _init_linear(linear: nn.Linear, weight_init: str = "xavier_uniform", bias_init: str = "zeros") -> None:
+    r"""
+    Initialize a linear layer the way the official generator does.
+
+    Args:
+        linear (`torch.nn.Linear`):
+            Layer to initialize in place.
+        weight_init (`str`, defaults to `"xavier_uniform"`):
+            `"xavier_uniform"`, `"zeros"`, or `"normal"` (std 0.02).
+        bias_init (`str`, defaults to `"zeros"`):
+            Bias fill. Only `"zeros"` is used by the released architecture.
+    """
+    if weight_init == "xavier_uniform":
+        nn.init.xavier_uniform_(linear.weight)
+    elif weight_init == "zeros":
+        nn.init.zeros_(linear.weight)
+    elif weight_init == "normal":
+        nn.init.normal_(linear.weight, std=0.02)
+    else:
+        raise ValueError(f"weight_init must be xavier_uniform, zeros, or normal, got {weight_init!r}.")
+    if linear.bias is not None:
+        if bias_init != "zeros":
+            raise ValueError(f"bias_init must be 'zeros', got {bias_init!r}.")
+        nn.init.zeros_(linear.bias)
+
+
+class _TorchLinear(nn.Module):
+    r"""
+    Linear wrapper whose parameter names match the official `TorchLinear` (`*.linear.weight`).
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        weight_init: str = "xavier_uniform",
+        bias_init: str = "zeros",
+    ):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        _init_linear(self.linear, weight_init=weight_init, bias_init=bias_init)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Apply the wrapped linear layer.
+
+        Args:
+            x (`torch.Tensor`):
+                Input of shape `(..., in_features)`.
+
+        Returns:
+            `torch.Tensor`: Output of shape `(..., out_features)`.
+        """
+        return self.linear(x)
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(int(dim)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Root-mean-square normalize the last dimension.
+
+        Args:
+            x (`torch.Tensor`):
+                Input of shape `(..., dim)`.
+
+        Returns:
+            `torch.Tensor`: Normalized tensor, same shape and dtype as `x`.
+        """
+        input_dtype = x.dtype
+        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
+        normalized = x * torch.rsqrt(variance + self.eps)
+        return (normalized * self.weight).to(dtype=input_dtype)
+
+
+def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    r"""
+    Apply an adaLN-Zero shift and scale.
+
+    Args:
+        x (`torch.Tensor`):
+            Token tensor `(batch, tokens, hidden)`.
+        shift (`torch.Tensor`):
+            Per-sample shift `(batch, hidden)`.
+        scale (`torch.Tensor`):
+            Per-sample scale `(batch, hidden)`.
+
+    Returns:
+        `torch.Tensor`: `x * (1 + scale) + shift`, broadcast over tokens.
+    """
+    return x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def _sincos_1d(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    r"""
+    Sine-cosine embedding of 1-D positions.
+
+    Args:
+        embed_dim (`int`):
+            Even embedding width.
+        pos (`numpy.ndarray`):
+            Positions, flattened inside the function.
+
+    Returns:
+        `numpy.ndarray`: Array of shape `(positions, embed_dim)`.
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError(f"embed_dim must be even, got {embed_dim}.")
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega = 1.0 / 10000 ** (omega / (embed_dim / 2.0))
+    out = np.einsum("m,d->md", pos.reshape(-1), omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+def _sincos_2d(embed_dim: int, grid_size: int) -> np.ndarray:
+    r"""
+    Sine-cosine embedding of a square grid, matching the official generator.
+
+    Args:
+        embed_dim (`int`):
+            Full embedding width. Half is used for each axis.
+        grid_size (`int`):
+            Tokens along one side of the latent grid.
+
+    Returns:
+        `numpy.ndarray`: Array of shape `(grid_size * grid_size, embed_dim)`.
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0).reshape(2, 1, grid_size, grid_size)
+    half = embed_dim // 2
+    return np.concatenate([_sincos_1d(half, grid[0]), _sincos_1d(half, grid[1])], axis=1)
+
+
+def _swiglu_hidden(hidden_size: int, mlp_ratio: float) -> int:
+    r"""
+    SwiGLU intermediate width used by the official blocks.
+
+    Args:
+        hidden_size (`int`):
+            Transformer width.
+        mlp_ratio (`float`):
+            Expansion ratio before the 2/3 SwiGLU reduction.
+
+    Returns:
+        `int`: Intermediate width, rounded up to a multiple of 32.
+    """
+    mlp_hidden = int(hidden_size * float(mlp_ratio))
+    hidden = int((2.0 / 3.0) * mlp_hidden)
+    return ((hidden + 31) // 32) * 32
+
+
+class _StandardMLP(nn.Module):
+    def __init__(self, hidden_size: int, mlp_hidden_dim: int):
+        super().__init__()
+        self.fc1 = _TorchLinear(hidden_size, mlp_hidden_dim, bias=True)
+        self.fc2 = _TorchLinear(mlp_hidden_dim, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        GELU MLP.
+
+        Args:
+            x (`torch.Tensor`):
+                Input `(..., hidden_size)`.
+
+        Returns:
+            `torch.Tensor`: Output `(..., hidden_size)`.
+        """
+        return self.fc2(nn.functional.gelu(self.fc1(x), approximate="none"))
+
+
+class _SwiGLUFFN(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.w1 = _TorchLinear(hidden_size, intermediate_size, bias=True)
+        self.w3 = _TorchLinear(hidden_size, intermediate_size, bias=True)
+        self.w2 = _TorchLinear(intermediate_size, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        SwiGLU feed-forward.
+
+        Args:
+            x (`torch.Tensor`):
+                Input `(..., hidden_size)`.
+
+        Returns:
+            `torch.Tensor`: Output `(..., hidden_size)`.
+        """
+        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
+
+
+class _Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        qk_norm: bool,
+        use_rmsnorm: bool,
+        use_rope: bool,
+        attn_fp32: bool,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.dim = int(dim)
+        self.num_heads = int(num_heads)
+        self.use_rope = bool(use_rope)
+        self.attn_fp32 = bool(attn_fp32)
+        self.qkv = _TorchLinear(self.dim, self.dim * 3, bias=True)
+        head_dim = self.dim // self.num_heads
+        if qk_norm:
+            if use_rmsnorm:
+                self.q_norm = _RMSNorm(head_dim)
+                self.k_norm = _RMSNorm(head_dim)
+            else:
+                self.q_norm = nn.LayerNorm(head_dim, eps=1e-6, elementwise_affine=True)
+                self.k_norm = nn.LayerNorm(head_dim, eps=1e-6, elementwise_affine=True)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+        self.proj = _TorchLinear(self.dim, self.dim, bias=True)
+        if use_rope:
+            half_dim = head_dim // 2
+            freqs = 1.0 / (10000 ** (torch.arange(0, half_dim, dtype=torch.float32) / half_dim))
+            positions = torch.arange(max_seq_len, dtype=torch.float32)
+            emb = torch.cat([torch.outer(positions, freqs)] * 2, dim=-1)
+            self.register_buffer("_rope_cos", torch.cos(emb)[None, :, None, :], persistent=False)
+            self.register_buffer("_rope_sin", torch.sin(emb)[None, :, None, :], persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Multi-head self-attention. Dropout is off, matching inference (`deterministic=True`).
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Attention output `(batch, tokens, hidden)`.
+        """
+        batch, tokens, _channels = x.shape
+        head_dim = self.dim // self.num_heads
+        qkv = self.qkv(x).reshape(batch, tokens, 3, self.num_heads, head_dim)
+        query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        if self.q_norm is not None:
+            query = self.q_norm(query)
+            key = self.k_norm(key)
+        if self.use_rope:
+            rope_dtype = torch.float32 if self.attn_fp32 else query.dtype
+            cos = self._rope_cos[:, :tokens].to(dtype=rope_dtype)
+            sin = self._rope_sin[:, :tokens].to(dtype=rope_dtype)
+            half_dim = head_dim // 2
+
+            def _rotate(tensor: torch.Tensor) -> torch.Tensor:
+                first, second = tensor[..., :half_dim], tensor[..., half_dim:]
+                return torch.cat([-second, first], dim=-1)
+
+            query = (query * cos) + (_rotate(query) * sin)
+            key = (key * cos) + (_rotate(key) * sin)
+        if self.attn_fp32:
+            query = query.float()
+            key = key.float()
+            value = value.float()
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+        if self.attn_fp32:
+            with torch.amp.autocast(device_type=query.device.type, enabled=False):
+                out = nn.functional.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        else:
+            out = nn.functional.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        out = out.permute(0, 2, 1, 3).reshape(batch, tokens, self.dim)
+        return self.proj(out)
+
+
+class _DiTBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float,
+        use_qknorm: bool,
+        use_swiglu: bool,
+        use_rmsnorm: bool,
+        use_rope: bool,
+        attn_fp32: bool,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.hidden_size = int(hidden_size)
+        if use_rmsnorm:
+            self.norm1 = _RMSNorm(self.hidden_size)
+            self.norm2 = _RMSNorm(self.hidden_size)
+        else:
+            self.norm1 = nn.LayerNorm(self.hidden_size, eps=1e-6, elementwise_affine=False)
+            self.norm2 = nn.LayerNorm(self.hidden_size, eps=1e-6, elementwise_affine=False)
+        self.attn = _Attention(
+            dim=self.hidden_size,
+            num_heads=num_heads,
+            qk_norm=use_qknorm,
+            use_rmsnorm=use_rmsnorm,
+            use_rope=use_rope,
+            attn_fp32=attn_fp32,
+            max_seq_len=max_seq_len,
+        )
+        if use_swiglu:
+            self.mlp = _SwiGLUFFN(self.hidden_size, _swiglu_hidden(self.hidden_size, mlp_ratio))
+        else:
+            self.mlp = _StandardMLP(self.hidden_size, int(self.hidden_size * float(mlp_ratio)))
+        self.adaLN_mod = nn.Sequential(
+            nn.SiLU(),
+            _TorchLinear(self.hidden_size, 6 * self.hidden_size, bias=True, weight_init="zeros"),
+        )
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        One Lightning-DiT block conditioned by adaLN-Zero.
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+            condition (`torch.Tensor`):
+                Per-sample condition `(batch, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Updated tokens, same shape as `x`.
+        """
+        chunks = self.adaLN_mod(condition.float()).to(dtype=x.dtype)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(chunks, 6, dim=1)
+        attended = self.attn(_modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + gate_msa.unsqueeze(1) * attended
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(_modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class _FinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int, use_rmsnorm: bool):
+        super().__init__()
+        self.norm_final = _RMSNorm(hidden_size) if use_rmsnorm else nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+        self.adaLN_mod = nn.Sequential(
+            nn.SiLU(),
+            _TorchLinear(hidden_size, 2 * hidden_size, bias=True, weight_init="zeros"),
+        )
+        self.linear = _TorchLinear(hidden_size, patch_size * patch_size * out_channels, bias=True, weight_init="zeros")
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        Project tokens to patch pixels.
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+            condition (`torch.Tensor`):
+                Per-sample condition `(batch, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Patch pixels `(batch, tokens, patch_size * patch_size * out_channels)`.
+        """
+        shift, scale = torch.chunk(self.adaLN_mod(condition.float()).to(dtype=x.dtype), 2, dim=1)
+        return self.linear(_modulate(self.norm_final(x), shift, scale))
+
+
+class _LightningDiT(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        patch_size: int,
+        in_channels: int,
+        hidden_size: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        out_channels: int,
+        use_qknorm: bool,
+        use_swiglu: bool,
+        use_rope: bool,
+        use_rmsnorm: bool,
+        cond_dim: int,
+        n_cls_tokens: int,
+        attn_fp32: bool,
+    ):
+        super().__init__()
+        self.input_size = int(input_size)
+        self.patch_size = int(patch_size)
+        self.in_channels = int(in_channels)
+        self.hidden_size = int(hidden_size)
+        self.out_channels = int(out_channels)
+        self.n_cls_tokens = int(n_cls_tokens)
+        self.patch_embed = _TorchLinear(self.patch_size * self.patch_size * self.in_channels, self.hidden_size, bias=True)
+        grid = self.input_size // self.patch_size
+        max_seq_len = grid * grid + self.n_cls_tokens
+        self.blocks = nn.ModuleList(
+            [
+                _DiTBlock(
+                    hidden_size=self.hidden_size,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    use_qknorm=use_qknorm,
+                    use_swiglu=use_swiglu,
+                    use_rmsnorm=use_rmsnorm,
+                    use_rope=use_rope,
+                    attn_fp32=attn_fp32,
+                    max_seq_len=max_seq_len,
+                )
+                for _ in range(int(depth))
+            ]
+        )
+        self.final_layer = _FinalLayer(self.hidden_size, self.patch_size, self.out_channels, use_rmsnorm)
+        position = _sincos_2d(self.hidden_size, grid)
+        self.pos_embed = nn.Parameter(torch.as_tensor(position, dtype=torch.float32).unsqueeze(0))
+        if self.n_cls_tokens > 0:
+            self.cls_embed = nn.Parameter(torch.randn(1, self.n_cls_tokens, self.hidden_size) * 0.02)
+            self.c_token_proj = _TorchLinear(int(cond_dim), self.hidden_size, bias=True)
+        else:
+            self.register_parameter("cls_embed", None)
+            self.c_token_proj = None
+
+    def forward(self, hidden_states_bhwc: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        Map a latent grid to a clean latent grid.
+
+        Args:
+            hidden_states_bhwc (`torch.Tensor`):
+                Reference noise `(batch, height, width, channels)`. Height and width must equal
+                `input_size`. The official generator uses this layout.
+            condition (`torch.Tensor`):
+                Class, noise-code, and guidance embedding `(batch, cond_dim)`.
+
+        Returns:
+            `torch.Tensor`: Clean latent `(batch, height, width, channels)`.
+        """
+        batch, height, width, channels = hidden_states_bhwc.shape
+        patch = self.patch_size
+        grid = self.input_size // patch
+        if (height, width, channels) != (self.input_size, self.input_size, self.in_channels):
+            raise ValueError(
+                f"Expected latent {(self.input_size, self.input_size, self.in_channels)} in BHWC, "
+                f"got {(height, width, channels)}."
+            )
+        tokens = hidden_states_bhwc.reshape(batch, grid, patch, grid, patch, channels)
+        tokens = tokens.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch, grid * grid, patch * patch * channels)
+        tokens = self.patch_embed(tokens)
+        tokens = (tokens + self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)).to(dtype=tokens.dtype)
+        if self.n_cls_tokens > 0:
+            class_tokens = self.c_token_proj(condition).unsqueeze(1).repeat(1, self.n_cls_tokens, 1)
+            class_tokens = class_tokens + self.cls_embed.to(device=tokens.device, dtype=tokens.dtype)
+            tokens = torch.cat([class_tokens, tokens], dim=1)
+        for block in self.blocks:
+            tokens = block(tokens, condition)
+        tokens = self.final_layer(tokens, condition)
+        if self.n_cls_tokens > 0:
+            tokens = tokens[:, self.n_cls_tokens :, :]
+        tokens = tokens.reshape(batch, grid, grid, patch, patch, self.out_channels)
+        tokens = tokens.permute(0, 1, 3, 2, 4, 5).contiguous()
+        return tokens.reshape(batch, self.input_size, self.input_size, self.out_channels)
+
+
+class _TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.frequency_embedding_size = int(frequency_embedding_size)
+        self.fc1 = _TorchLinear(self.frequency_embedding_size, hidden_size, bias=True, weight_init="normal")
+        self.fc2 = _TorchLinear(hidden_size, hidden_size, bias=True, weight_init="normal")
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        r"""
+        Embed a scalar guidance value with sinusoidal features.
+
+        Args:
+            timesteps (`torch.Tensor`):
+                Shape `(batch,)`. This is the guidance scale, not a diffusion time.
+
+        Returns:
+            `torch.Tensor`: Embedding `(batch, hidden)`.
+        """
+        timesteps = timesteps.to(dtype=torch.float32)
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(0, half, device=timesteps.device, dtype=torch.float32) / half)
+        arguments = timesteps[:, None] * freqs[None]
+        embedding = torch.cat([torch.cos(arguments), torch.sin(arguments)], dim=-1)
+        if self.frequency_embedding_size % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        hidden = nn.functional.silu(self.fc1(embedding))
+        return self.fc2(hidden)
+
+
+class WFlowTransformer2DModel(ModelMixin, ConfigMixin):
+    r"""
+    One-step ImageNet generator used by W-Flow.
+
+    The network maps unit Gaussian SD-VAE noise plus a class label to a clean latent in one pass.
+    Guidance is a continuous embedding (the official `cfg_scale`, equal to `w + 1` in the paper),
+    not a second unconditional forward. Discrete noise codes (`noise_classes` by `noise_coords`)
+    are added to the class embedding.
+
+    Parameter names match `DitGen` in the official repository, so an `ema_model` state dict loads
+    without renaming.
+
+    Args:
+        sample_size (`int`, defaults to 32):
+            Latent spatial size. Released checkpoints use 32, which decodes to 256 pixels.
+        patch_size (`int`, defaults to 2):
+            Patch size on the latent grid.
+        in_channels (`int`, defaults to 4):
+            Latent channels.
+        out_channels (`int`, defaults to 4):
+            Predicted latent channels.
+        cond_dim (`int`, defaults to 768):
+            Width of the class, noise-code, and guidance embeddings.
+        num_classes (`int`, defaults to 1000):
+            ImageNet class count. There is no null-class slot.
+        noise_classes (`int`, defaults to 64):
+            Cardinality of each discrete noise code.
+        noise_coords (`int`, defaults to 32):
+            Number of independent noise-code embeddings that are summed.
+        n_cls_tokens (`int`, defaults to 16):
+            Learned class tokens prepended to the patch sequence.
+        hidden_size (`int`, defaults to 768):
+            Transformer width.
+        depth (`int`, defaults to 12):
+            Number of blocks.
+        num_attention_heads (`int`, defaults to 12):
+            Attention heads.
+        mlp_ratio (`float`, defaults to 4.0):
+            MLP expansion before the SwiGLU 2/3 reduction.
+        use_qknorm (`bool`, defaults to `True`):
+            QK normalization.
+        use_swiglu (`bool`, defaults to `True`):
+            SwiGLU MLP. Released checkpoints turn this on.
+        use_rope (`bool`, defaults to `True`):
+            1-D rotary embeddings over the token sequence, including class tokens.
+        use_rmsnorm (`bool`, defaults to `True`):
+            RMSNorm instead of LayerNorm.
+        use_bf16 (`bool`, defaults to `True`):
+            Autocast the forward to bfloat16 on CUDA, matching the official sampler.
+        attn_fp32 (`bool`, defaults to `False`):
+            Run attention in float32. The ablation checkpoint sets this to `True`.
+        guidance_scale (`float`, defaults to 1.0):
+            Default embedded guidance. The pipeline reads this when the caller omits a scale.
+        default_num_inference_steps (`int`, defaults to 1):
+            Sampling is one network evaluation.
+        model_type (`str`, defaults to `"W-Flow-B"`):
+            Variant name stored in the config.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        cond_dim: int = 768,
+        num_classes: int = 1000,
+        noise_classes: int = 64,
+        noise_coords: int = 32,
+        n_cls_tokens: int = 16,
+        hidden_size: int = 768,
+        depth: int = 12,
+        num_attention_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        use_qknorm: bool = True,
+        use_swiglu: bool = True,
+        use_rope: bool = True,
+        use_rmsnorm: bool = True,
+        use_bf16: bool = True,
+        attn_fp32: bool = False,
+        guidance_scale: float = 1.0,
+        default_num_inference_steps: int = 1,
+        model_type: str = "W-Flow-B",
+    ):
+        super().__init__()
+        self.class_embed = nn.Embedding(int(num_classes), int(cond_dim))
+        nn.init.normal_(self.class_embed.weight, std=0.02)
+        self.noise_embeds = nn.ModuleList()
+        for _ in range(int(noise_coords)):
+            embedding = nn.Embedding(int(noise_classes), int(cond_dim))
+            nn.init.normal_(embedding.weight, std=0.02)
+            self.noise_embeds.append(embedding)
+        self.cfg_embedder = _TimestepEmbedder(int(cond_dim))
+        self.cfg_norm = _RMSNorm(int(cond_dim))
+        self.model = _LightningDiT(
+            input_size=int(sample_size),
+            patch_size=int(patch_size),
+            in_channels=int(in_channels),
+            hidden_size=int(hidden_size),
+            depth=int(depth),
+            num_heads=int(num_attention_heads),
+            mlp_ratio=float(mlp_ratio),
+            out_channels=int(out_channels),
+            use_qknorm=bool(use_qknorm),
+            use_swiglu=bool(use_swiglu),
+            use_rope=bool(use_rope),
+            use_rmsnorm=bool(use_rmsnorm),
+            cond_dim=int(cond_dim),
+            n_cls_tokens=int(n_cls_tokens),
+            attn_fp32=bool(attn_fp32),
+        )
+
+    def condition(
+        self,
+        class_labels: torch.Tensor,
+        guidance_scale: float | torch.Tensor,
+        noise_labels: torch.Tensor,
+    ) -> torch.Tensor:
+        r"""
+        Build the conditioning vector from class, noise codes, and guidance.
+
+        Args:
+            class_labels (`torch.Tensor`):
+                Integer class ids `(batch,)`, in `[0, num_classes)`.
+            guidance_scale (`float` or `torch.Tensor`):
+                Official CFG value `(w + 1)`. A float is broadcast across the batch. A tensor
+                must be `(batch,)` or a scalar.
+            noise_labels (`torch.Tensor`):
+                Integer codes `(batch, noise_coords)` in `[0, noise_classes)`.
+
+        Returns:
+            `torch.Tensor`: Condition `(batch, cond_dim)`.
+        """
+        batch = class_labels.shape[0]
+        cond = self.class_embed(class_labels)
+        for index, embedding in enumerate(self.noise_embeds):
+            cond = cond + embedding(noise_labels[:, index])
+        if isinstance(guidance_scale, (float, int)):
+            scale = torch.full((batch,), float(guidance_scale), device=class_labels.device, dtype=torch.float32)
+        else:
+            scale = torch.as_tensor(guidance_scale, device=class_labels.device, dtype=torch.float32)
+            if scale.ndim == 0:
+                scale = scale[None].repeat(batch)
+        embedded = self.cfg_norm(self.cfg_embedder(scale.float()).float())
+        return cond + embedded.to(dtype=cond.dtype) * 0.02
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        class_labels: torch.Tensor,
+        guidance_scale: float | torch.Tensor = 1.0,
+        noise_labels: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ) -> WFlowTransformer2DModelOutput | tuple[torch.Tensor]:
+        r"""
+        Predict a clean latent from reference noise.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                Gaussian reference noise `(batch, channels, height, width)` in BCHW. Height and
+                width must equal `config.sample_size`.
+            class_labels (`torch.Tensor`):
+                Integer ImageNet ids `(batch,)`.
+            guidance_scale (`float` or `torch.Tensor`, defaults to 1.0):
+                Embedded guidance. This is the official `cfg_scale` (`w + 1` in the paper).
+            noise_labels (`torch.Tensor`, *optional*):
+                Discrete noise codes `(batch, noise_coords)`. Required. The pipeline samples these
+                with the same generator as `hidden_states`.
+            return_dict (`bool`, defaults to `True`):
+                Return [`WFlowTransformer2DModelOutput`] when `True`, otherwise a one-tuple.
+
+        Returns:
+            [`WFlowTransformer2DModelOutput`] or `tuple`:
+                `sample` is the clean latent `(batch, channels, height, width)`.
+        """
+        if noise_labels is None:
+            raise ValueError("`noise_labels` is required and must have shape (batch, noise_coords).")
+        if hidden_states.ndim != 4:
+            raise ValueError(f"`hidden_states` must be BCHW, got {tuple(hidden_states.shape)}.")
+        hidden_bhwc = hidden_states.permute(0, 2, 3, 1).contiguous()
+        use_autocast = bool(self.config.use_bf16) and hidden_states.device.type == "cuda"
+        with torch.amp.autocast(device_type=hidden_states.device.type, dtype=torch.bfloat16, enabled=use_autocast):
+            cond = self.condition(class_labels, guidance_scale, noise_labels)
+            sample_bhwc = self.model(hidden_bhwc, cond)
+        sample = sample_bhwc.permute(0, 3, 1, 2).contiguous()
+        if not return_dict:
+            return (sample,)
+        return WFlowTransformer2DModelOutput(sample=sample)

W-Flow-B/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

W-Flow-B/vae/diffusion_pytorch_model.safetensors

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

W-Flow-L/model_index.json

@@ -0,0 +1,1021 @@
+{
+  "_class_name": [
+    "pipeline",
+    "WFlowPipeline"
+  ],
+  "_diffusers_version": "0.39.0",
+  "scheduler": [
+    "scheduling_wflow",
+    "WFlowScheduler"
+  ],
+  "transformer": [
+    "transformer_wflow",
+    "WFlowTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ],
+  "id2label": {
+    "0": "tench, Tinca tinca",
+    "1": "goldfish, Carassius auratus",
+    "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
+    "3": "tiger shark, Galeocerdo cuvieri",
+    "4": "hammerhead, hammerhead shark",
+    "5": "electric ray, crampfish, numbfish, torpedo",
+    "6": "stingray",
+    "7": "cock",
+    "8": "hen",
+    "9": "ostrich, Struthio camelus",
+    "10": "brambling, Fringilla montifringilla",
+    "11": "goldfinch, Carduelis carduelis",
+    "12": "house finch, linnet, Carpodacus mexicanus",
+    "13": "junco, snowbird",
+    "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+    "15": "robin, American robin, Turdus migratorius",
+    "16": "bulbul",
+    "17": "jay",
+    "18": "magpie",
+    "19": "chickadee",
+    "20": "water ouzel, dipper",
+    "21": "kite",
+    "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
+    "23": "vulture",
+    "24": "great grey owl, great gray owl, Strix nebulosa",
+    "25": "European fire salamander, Salamandra salamandra",
+    "26": "common newt, Triturus vulgaris",
+    "27": "eft",
+    "28": "spotted salamander, Ambystoma maculatum",
+    "29": "axolotl, mud puppy, Ambystoma mexicanum",
+    "30": "bullfrog, Rana catesbeiana",
+    "31": "tree frog, tree-frog",
+    "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+    "33": "loggerhead, loggerhead turtle, Caretta caretta",
+    "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
+    "35": "mud turtle",
+    "36": "terrapin",
+    "37": "box turtle, box tortoise",
+    "38": "banded gecko",
+    "39": "common iguana, iguana, Iguana iguana",
+    "40": "American chameleon, anole, Anolis carolinensis",
+    "41": "whiptail, whiptail lizard",
+    "42": "agama",
+    "43": "frilled lizard, Chlamydosaurus kingi",
+    "44": "alligator lizard",
+    "45": "Gila monster, Heloderma suspectum",
+    "46": "green lizard, Lacerta viridis",
+    "47": "African chameleon, Chamaeleo chamaeleon",
+    "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
+    "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
+    "50": "American alligator, Alligator mississipiensis",
+    "51": "triceratops",
+    "52": "thunder snake, worm snake, Carphophis amoenus",
+    "53": "ringneck snake, ring-necked snake, ring snake",
+    "54": "hognose snake, puff adder, sand viper",
+    "55": "green snake, grass snake",
+    "56": "king snake, kingsnake",
+    "57": "garter snake, grass snake",
+    "58": "water snake",
+    "59": "vine snake",
+    "60": "night snake, Hypsiglena torquata",
+    "61": "boa constrictor, Constrictor constrictor",
+    "62": "rock python, rock snake, Python sebae",
+    "63": "Indian cobra, Naja naja",
+    "64": "green mamba",
+    "65": "sea snake",
+    "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+    "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+    "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
+    "69": "trilobite",
+    "70": "harvestman, daddy longlegs, Phalangium opilio",
+    "71": "scorpion",
+    "72": "black and gold garden spider, Argiope aurantia",
+    "73": "barn spider, Araneus cavaticus",
+    "74": "garden spider, Aranea diademata",
+    "75": "black widow, Latrodectus mactans",
+    "76": "tarantula",
+    "77": "wolf spider, hunting spider",
+    "78": "tick",
+    "79": "centipede",
+    "80": "black grouse",
+    "81": "ptarmigan",
+    "82": "ruffed grouse, partridge, Bonasa umbellus",
+    "83": "prairie chicken, prairie grouse, prairie fowl",
+    "84": "peacock",
+    "85": "quail",
+    "86": "partridge",
+    "87": "African grey, African gray, Psittacus erithacus",
+    "88": "macaw",
+    "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+    "90": "lorikeet",
+    "91": "coucal",
+    "92": "bee eater",
+    "93": "hornbill",
+    "94": "hummingbird",
+    "95": "jacamar",
+    "96": "toucan",
+    "97": "drake",
+    "98": "red-breasted merganser, Mergus serrator",
+    "99": "goose",
+    "100": "black swan, Cygnus atratus",
+    "101": "tusker",
+    "102": "echidna, spiny anteater, anteater",
+    "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",
+    "104": "wallaby, brush kangaroo",
+    "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",
+    "106": "wombat",
+    "107": "jellyfish",
+    "108": "sea anemone, anemone",
+    "109": "brain coral",
+    "110": "flatworm, platyhelminth",
+    "111": "nematode, nematode worm, roundworm",
+    "112": "conch",
+    "113": "snail",
+    "114": "slug",
+    "115": "sea slug, nudibranch",
+    "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore",
+    "117": "chambered nautilus, pearly nautilus, nautilus",
+    "118": "Dungeness crab, Cancer magister",
+    "119": "rock crab, Cancer irroratus",
+    "120": "fiddler crab",
+    "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica",
+    "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus",
+    "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish",
+    "124": "crayfish, crawfish, crawdad, crawdaddy",
+    "125": "hermit crab",
+    "126": "isopod",
+    "127": "white stork, Ciconia ciconia",
+    "128": "black stork, Ciconia nigra",
+    "129": "spoonbill",
+    "130": "flamingo",
+    "131": "little blue heron, Egretta caerulea",
+    "132": "American egret, great white heron, Egretta albus",
+    "133": "bittern",
+    "134": "crane",
+    "135": "limpkin, Aramus pictus",
+    "136": "European gallinule, Porphyrio porphyrio",
+    "137": "American coot, marsh hen, mud hen, water hen, Fulica americana",
+    "138": "bustard",
+    "139": "ruddy turnstone, Arenaria interpres",
+    "140": "red-backed sandpiper, dunlin, Erolia alpina",
+    "141": "redshank, Tringa totanus",
+    "142": "dowitcher",
+    "143": "oystercatcher, oyster catcher",
+    "144": "pelican",
+    "145": "king penguin, Aptenodytes patagonica",
+    "146": "albatross, mollymawk",
+    "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus",
+    "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca",
+    "149": "dugong, Dugong dugon",
+    "150": "sea lion",
+    "151": "Chihuahua",
+    "152": "Japanese spaniel",
+    "153": "Maltese dog, Maltese terrier, Maltese",
+    "154": "Pekinese, Pekingese, Peke",
+    "155": "Shih-Tzu",
+    "156": "Blenheim spaniel",
+    "157": "papillon",
+    "158": "toy terrier",
+    "159": "Rhodesian ridgeback",
+    "160": "Afghan hound, Afghan",
+    "161": "basset, basset hound",
+    "162": "beagle",
+    "163": "bloodhound, sleuthhound",
+    "164": "bluetick",
+    "165": "black-and-tan coonhound",
+    "166": "Walker hound, Walker foxhound",
+    "167": "English foxhound",
+    "168": "redbone",
+    "169": "borzoi, Russian wolfhound",
+    "170": "Irish wolfhound",
+    "171": "Italian greyhound",
+    "172": "whippet",
+    "173": "Ibizan hound, Ibizan Podenco",
+    "174": "Norwegian elkhound, elkhound",
+    "175": "otterhound, otter hound",
+    "176": "Saluki, gazelle hound",
+    "177": "Scottish deerhound, deerhound",
+    "178": "Weimaraner",
+    "179": "Staffordshire bullterrier, Staffordshire bull terrier",
+    "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",
+    "181": "Bedlington terrier",
+    "182": "Border terrier",
+    "183": "Kerry blue terrier",
+    "184": "Irish terrier",
+    "185": "Norfolk terrier",
+    "186": "Norwich terrier",
+    "187": "Yorkshire terrier",
+    "188": "wire-haired fox terrier",
+    "189": "Lakeland terrier",
+    "190": "Sealyham terrier, Sealyham",
+    "191": "Airedale, Airedale terrier",
+    "192": "cairn, cairn terrier",
+    "193": "Australian terrier",
+    "194": "Dandie Dinmont, Dandie Dinmont terrier",
+    "195": "Boston bull, Boston terrier",
+    "196": "miniature schnauzer",
+    "197": "giant schnauzer",
+    "198": "standard schnauzer",
+    "199": "Scotch terrier, Scottish terrier, Scottie",
+    "200": "Tibetan terrier, chrysanthemum dog",
+    "201": "silky terrier, Sydney silky",
+    "202": "soft-coated wheaten terrier",
+    "203": "West Highland white terrier",
+    "204": "Lhasa, Lhasa apso",
+    "205": "flat-coated retriever",
+    "206": "curly-coated retriever",
+    "207": "golden retriever",
+    "208": "Labrador retriever",
+    "209": "Chesapeake Bay retriever",
+    "210": "German short-haired pointer",
+    "211": "vizsla, Hungarian pointer",
+    "212": "English setter",
+    "213": "Irish setter, red setter",
+    "214": "Gordon setter",
+    "215": "Brittany spaniel",
+    "216": "clumber, clumber spaniel",
+    "217": "English springer, English springer spaniel",
+    "218": "Welsh springer spaniel",
+    "219": "cocker spaniel, English cocker spaniel, cocker",
+    "220": "Sussex spaniel",
+    "221": "Irish water spaniel",
+    "222": "kuvasz",
+    "223": "schipperke",
+    "224": "groenendael",
+    "225": "malinois",
+    "226": "briard",
+    "227": "kelpie",
+    "228": "komondor",
+    "229": "Old English sheepdog, bobtail",
+    "230": "Shetland sheepdog, Shetland sheep dog, Shetland",
+    "231": "collie",
+    "232": "Border collie",
+    "233": "Bouvier des Flandres, Bouviers des Flandres",
+    "234": "Rottweiler",
+    "235": "German shepherd, German shepherd dog, German police dog, alsatian",
+    "236": "Doberman, Doberman pinscher",
+    "237": "miniature pinscher",
+    "238": "Greater Swiss Mountain dog",
+    "239": "Bernese mountain dog",
+    "240": "Appenzeller",
+    "241": "EntleBucher",
+    "242": "boxer",
+    "243": "bull mastiff",
+    "244": "Tibetan mastiff",
+    "245": "French bulldog",
+    "246": "Great Dane",
+    "247": "Saint Bernard, St Bernard",
+    "248": "Eskimo dog, husky",
+    "249": "malamute, malemute, Alaskan malamute",
+    "250": "Siberian husky",
+    "251": "dalmatian, coach dog, carriage dog",
+    "252": "affenpinscher, monkey pinscher, monkey dog",
+    "253": "basenji",
+    "254": "pug, pug-dog",
+    "255": "Leonberg",
+    "256": "Newfoundland, Newfoundland dog",
+    "257": "Great Pyrenees",
+    "258": "Samoyed, Samoyede",
+    "259": "Pomeranian",
+    "260": "chow, chow chow",
+    "261": "keeshond",
+    "262": "Brabancon griffon",
+    "263": "Pembroke, Pembroke Welsh corgi",
+    "264": "Cardigan, Cardigan Welsh corgi",
+    "265": "toy poodle",
+    "266": "miniature poodle",
+    "267": "standard poodle",
+    "268": "Mexican hairless",
+    "269": "timber wolf, grey wolf, gray wolf, Canis lupus",
+    "270": "white wolf, Arctic wolf, Canis lupus tundrarum",
+    "271": "red wolf, maned wolf, Canis rufus, Canis niger",
+    "272": "coyote, prairie wolf, brush wolf, Canis latrans",
+    "273": "dingo, warrigal, warragal, Canis dingo",
+    "274": "dhole, Cuon alpinus",
+    "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
+    "276": "hyena, hyaena",
+    "277": "red fox, Vulpes vulpes",
+    "278": "kit fox, Vulpes macrotis",
+    "279": "Arctic fox, white fox, Alopex lagopus",
+    "280": "grey fox, gray fox, Urocyon cinereoargenteus",
+    "281": "tabby, tabby cat",
+    "282": "tiger cat",
+    "283": "Persian cat",
+    "284": "Siamese cat, Siamese",
+    "285": "Egyptian cat",
+    "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor",
+    "287": "lynx, catamount",
+    "288": "leopard, Panthera pardus",
+    "289": "snow leopard, ounce, Panthera uncia",
+    "290": "jaguar, panther, Panthera onca, Felis onca",
+    "291": "lion, king of beasts, Panthera leo",
+    "292": "tiger, Panthera tigris",
+    "293": "cheetah, chetah, Acinonyx jubatus",
+    "294": "brown bear, bruin, Ursus arctos",
+    "295": "American black bear, black bear, Ursus americanus, Euarctos americanus",
+    "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus",
+    "297": "sloth bear, Melursus ursinus, Ursus ursinus",
+    "298": "mongoose",
+    "299": "meerkat, mierkat",
+    "300": "tiger beetle",
+    "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle",
+    "302": "ground beetle, carabid beetle",
+    "303": "long-horned beetle, longicorn, longicorn beetle",
+    "304": "leaf beetle, chrysomelid",
+    "305": "dung beetle",
+    "306": "rhinoceros beetle",
+    "307": "weevil",
+    "308": "fly",
+    "309": "bee",
+    "310": "ant, emmet, pismire",
+    "311": "grasshopper, hopper",
+    "312": "cricket",
+    "313": "walking stick, walkingstick, stick insect",
+    "314": "cockroach, roach",
+    "315": "mantis, mantid",
+    "316": "cicada, cicala",
+    "317": "leafhopper",
+    "318": "lacewing, lacewing fly",
+    "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",
+    "320": "damselfly",
+    "321": "admiral",
+    "322": "ringlet, ringlet butterfly",
+    "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus",
+    "324": "cabbage butterfly",
+    "325": "sulphur butterfly, sulfur butterfly",
+    "326": "lycaenid, lycaenid butterfly",
+    "327": "starfish, sea star",
+    "328": "sea urchin",
+    "329": "sea cucumber, holothurian",
+    "330": "wood rabbit, cottontail, cottontail rabbit",
+    "331": "hare",
+    "332": "Angora, Angora rabbit",
+    "333": "hamster",
+    "334": "porcupine, hedgehog",
+    "335": "fox squirrel, eastern fox squirrel, Sciurus niger",
+    "336": "marmot",
+    "337": "beaver",
+    "338": "guinea pig, Cavia cobaya",
+    "339": "sorrel",
+    "340": "zebra",
+    "341": "hog, pig, grunter, squealer, Sus scrofa",
+    "342": "wild boar, boar, Sus scrofa",
+    "343": "warthog",
+    "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius",
+    "345": "ox",
+    "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis",
+    "347": "bison",
+    "348": "ram, tup",
+    "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis",
+    "350": "ibex, Capra ibex",
+    "351": "hartebeest",
+    "352": "impala, Aepyceros melampus",
+    "353": "gazelle",
+    "354": "Arabian camel, dromedary, Camelus dromedarius",
+    "355": "llama",
+    "356": "weasel",
+    "357": "mink",
+    "358": "polecat, fitch, foulmart, foumart, Mustela putorius",
+    "359": "black-footed ferret, ferret, Mustela nigripes",
+    "360": "otter",
+    "361": "skunk, polecat, wood pussy",
+    "362": "badger",
+    "363": "armadillo",
+    "364": "three-toed sloth, ai, Bradypus tridactylus",
+    "365": "orangutan, orang, orangutang, Pongo pygmaeus",
+    "366": "gorilla, Gorilla gorilla",
+    "367": "chimpanzee, chimp, Pan troglodytes",
+    "368": "gibbon, Hylobates lar",
+    "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus",
+    "370": "guenon, guenon monkey",
+    "371": "patas, hussar monkey, Erythrocebus patas",
+    "372": "baboon",
+    "373": "macaque",
+    "374": "langur",
+    "375": "colobus, colobus monkey",
+    "376": "proboscis monkey, Nasalis larvatus",
+    "377": "marmoset",
+    "378": "capuchin, ringtail, Cebus capucinus",
+    "379": "howler monkey, howler",
+    "380": "titi, titi monkey",
+    "381": "spider monkey, Ateles geoffroyi",
+    "382": "squirrel monkey, Saimiri sciureus",
+    "383": "Madagascar cat, ring-tailed lemur, Lemur catta",
+    "384": "indri, indris, Indri indri, Indri brevicaudatus",
+    "385": "Indian elephant, Elephas maximus",
+    "386": "African elephant, Loxodonta africana",
+    "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens",
+    "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca",
+    "389": "barracouta, snoek",
+    "390": "eel",
+    "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch",
+    "392": "rock beauty, Holocanthus tricolor",
+    "393": "anemone fish",
+    "394": "sturgeon",
+    "395": "gar, garfish, garpike, billfish, Lepisosteus osseus",
+    "396": "lionfish",
+    "397": "puffer, pufferfish, blowfish, globefish",
+    "398": "abacus",
+    "399": "abaya",
+    "400": "academic gown, academic robe, judge robe",
+    "401": "accordion, piano accordion, squeeze box",
+    "402": "acoustic guitar",
+    "403": "aircraft carrier, carrier, flattop, attack aircraft carrier",
+    "404": "airliner",
+    "405": "airship, dirigible",
+    "406": "altar",
+    "407": "ambulance",
+    "408": "amphibian, amphibious vehicle",
+    "409": "analog clock",
+    "410": "apiary, bee house",
+    "411": "apron",
+    "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin",
+    "413": "assault rifle, assault gun",
+    "414": "backpack, back pack, knapsack, packsack, rucksack, haversack",
+    "415": "bakery, bakeshop, bakehouse",
+    "416": "balance beam, beam",
+    "417": "balloon",
+    "418": "ballpoint, ballpoint pen, ballpen, Biro",
+    "419": "Band Aid",
+    "420": "banjo",
+    "421": "bannister, banister, balustrade, balusters, handrail",
+    "422": "barbell",
+    "423": "barber chair",
+    "424": "barbershop",
+    "425": "barn",
+    "426": "barometer",
+    "427": "barrel, cask",
+    "428": "barrow, garden cart, lawn cart, wheelbarrow",
+    "429": "baseball",
+    "430": "basketball",
+    "431": "bassinet",
+    "432": "bassoon",
+    "433": "bathing cap, swimming cap",
+    "434": "bath towel",
+    "435": "bathtub, bathing tub, bath, tub",
+    "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon",
+    "437": "beacon, lighthouse, beacon light, pharos",
+    "438": "beaker",
+    "439": "bearskin, busby, shako",
+    "440": "beer bottle",
+    "441": "beer glass",
+    "442": "bell cote, bell cot",
+    "443": "bib",
+    "444": "bicycle-built-for-two, tandem bicycle, tandem",
+    "445": "bikini, two-piece",
+    "446": "binder, ring-binder",
+    "447": "binoculars, field glasses, opera glasses",
+    "448": "birdhouse",
+    "449": "boathouse",
+    "450": "bobsled, bobsleigh, bob",
+    "451": "bolo tie, bolo, bola tie, bola",
+    "452": "bonnet, poke bonnet",
+    "453": "bookcase",
+    "454": "bookshop, bookstore, bookstall",
+    "455": "bottlecap",
+    "456": "bow",
+    "457": "bow tie, bow-tie, bowtie",
+    "458": "brass, memorial tablet, plaque",
+    "459": "brassiere, bra, bandeau",
+    "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty",
+    "461": "breastplate, aegis, egis",
+    "462": "broom",
+    "463": "bucket, pail",
+    "464": "buckle",
+    "465": "bulletproof vest",
+    "466": "bullet train, bullet",
+    "467": "butcher shop, meat market",
+    "468": "cab, hack, taxi, taxicab",
+    "469": "caldron, cauldron",
+    "470": "candle, taper, wax light",
+    "471": "cannon",
+    "472": "canoe",
+    "473": "can opener, tin opener",
+    "474": "cardigan",
+    "475": "car mirror",
+    "476": "carousel, carrousel, merry-go-round, roundabout, whirligig",
+    "477": "carpenters kit, tool kit",
+    "478": "carton",
+    "479": "car wheel",
+    "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM",
+    "481": "cassette",
+    "482": "cassette player",
+    "483": "castle",
+    "484": "catamaran",
+    "485": "CD player",
+    "486": "cello, violoncello",
+    "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone",
+    "488": "chain",
+    "489": "chainlink fence",
+    "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour",
+    "491": "chain saw, chainsaw",
+    "492": "chest",
+    "493": "chiffonier, commode",
+    "494": "chime, bell, gong",
+    "495": "china cabinet, china closet",
+    "496": "Christmas stocking",
+    "497": "church, church building",
+    "498": "cinema, movie theater, movie theatre, movie house, picture palace",
+    "499": "cleaver, meat cleaver, chopper",
+    "500": "cliff dwelling",
+    "501": "cloak",
+    "502": "clog, geta, patten, sabot",
+    "503": "cocktail shaker",
+    "504": "coffee mug",
+    "505": "coffeepot",
+    "506": "coil, spiral, volute, whorl, helix",
+    "507": "combination lock",
+    "508": "computer keyboard, keypad",
+    "509": "confectionery, confectionary, candy store",
+    "510": "container ship, containership, container vessel",
+    "511": "convertible",
+    "512": "corkscrew, bottle screw",
+    "513": "cornet, horn, trumpet, trump",
+    "514": "cowboy boot",
+    "515": "cowboy hat, ten-gallon hat",
+    "516": "cradle",
+    "517": "crane",
+    "518": "crash helmet",
+    "519": "crate",
+    "520": "crib, cot",
+    "521": "Crock Pot",
+    "522": "croquet ball",
+    "523": "crutch",
+    "524": "cuirass",
+    "525": "dam, dike, dyke",
+    "526": "desk",
+    "527": "desktop computer",
+    "528": "dial telephone, dial phone",
+    "529": "diaper, nappy, napkin",
+    "530": "digital clock",
+    "531": "digital watch",
+    "532": "dining table, board",
+    "533": "dishrag, dishcloth",
+    "534": "dishwasher, dish washer, dishwashing machine",
+    "535": "disk brake, disc brake",
+    "536": "dock, dockage, docking facility",
+    "537": "dogsled, dog sled, dog sleigh",
+    "538": "dome",
+    "539": "doormat, welcome mat",
+    "540": "drilling platform, offshore rig",
+    "541": "drum, membranophone, tympan",
+    "542": "drumstick",
+    "543": "dumbbell",
+    "544": "Dutch oven",
+    "545": "electric fan, blower",
+    "546": "electric guitar",
+    "547": "electric locomotive",
+    "548": "entertainment center",
+    "549": "envelope",
+    "550": "espresso maker",
+    "551": "face powder",
+    "552": "feather boa, boa",
+    "553": "file, file cabinet, filing cabinet",
+    "554": "fireboat",
+    "555": "fire engine, fire truck",
+    "556": "fire screen, fireguard",
+    "557": "flagpole, flagstaff",
+    "558": "flute, transverse flute",
+    "559": "folding chair",
+    "560": "football helmet",
+    "561": "forklift",
+    "562": "fountain",
+    "563": "fountain pen",
+    "564": "four-poster",
+    "565": "freight car",
+    "566": "French horn, horn",
+    "567": "frying pan, frypan, skillet",
+    "568": "fur coat",
+    "569": "garbage truck, dustcart",
+    "570": "gasmask, respirator, gas helmet",
+    "571": "gas pump, gasoline pump, petrol pump, island dispenser",
+    "572": "goblet",
+    "573": "go-kart",
+    "574": "golf ball",
+    "575": "golfcart, golf cart",
+    "576": "gondola",
+    "577": "gong, tam-tam",
+    "578": "gown",
+    "579": "grand piano, grand",
+    "580": "greenhouse, nursery, glasshouse",
+    "581": "grille, radiator grille",
+    "582": "grocery store, grocery, food market, market",
+    "583": "guillotine",
+    "584": "hair slide",
+    "585": "hair spray",
+    "586": "half track",
+    "587": "hammer",
+    "588": "hamper",
+    "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier",
+    "590": "hand-held computer, hand-held microcomputer",
+    "591": "handkerchief, hankie, hanky, hankey",
+    "592": "hard disc, hard disk, fixed disk",
+    "593": "harmonica, mouth organ, harp, mouth harp",
+    "594": "harp",
+    "595": "harvester, reaper",
+    "596": "hatchet",
+    "597": "holster",
+    "598": "home theater, home theatre",
+    "599": "honeycomb",
+    "600": "hook, claw",
+    "601": "hoopskirt, crinoline",
+    "602": "horizontal bar, high bar",
+    "603": "horse cart, horse-cart",
+    "604": "hourglass",
+    "605": "iPod",
+    "606": "iron, smoothing iron",
+    "607": "jack-o-lantern",
+    "608": "jean, blue jean, denim",
+    "609": "jeep, landrover",
+    "610": "jersey, T-shirt, tee shirt",
+    "611": "jigsaw puzzle",
+    "612": "jinrikisha, ricksha, rickshaw",
+    "613": "joystick",
+    "614": "kimono",
+    "615": "knee pad",
+    "616": "knot",
+    "617": "lab coat, laboratory coat",
+    "618": "ladle",
+    "619": "lampshade, lamp shade",
+    "620": "laptop, laptop computer",
+    "621": "lawn mower, mower",
+    "622": "lens cap, lens cover",
+    "623": "letter opener, paper knife, paperknife",
+    "624": "library",
+    "625": "lifeboat",
+    "626": "lighter, light, igniter, ignitor",
+    "627": "limousine, limo",
+    "628": "liner, ocean liner",
+    "629": "lipstick, lip rouge",
+    "630": "Loafer",
+    "631": "lotion",
+    "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system",
+    "633": "loupe, jewelers loupe",
+    "634": "lumbermill, sawmill",
+    "635": "magnetic compass",
+    "636": "mailbag, postbag",
+    "637": "mailbox, letter box",
+    "638": "maillot",
+    "639": "maillot, tank suit",
+    "640": "manhole cover",
+    "641": "maraca",
+    "642": "marimba, xylophone",
+    "643": "mask",
+    "644": "matchstick",
+    "645": "maypole",
+    "646": "maze, labyrinth",
+    "647": "measuring cup",
+    "648": "medicine chest, medicine cabinet",
+    "649": "megalith, megalithic structure",
+    "650": "microphone, mike",
+    "651": "microwave, microwave oven",
+    "652": "military uniform",
+    "653": "milk can",
+    "654": "minibus",
+    "655": "miniskirt, mini",
+    "656": "minivan",
+    "657": "missile",
+    "658": "mitten",
+    "659": "mixing bowl",
+    "660": "mobile home, manufactured home",
+    "661": "Model T",
+    "662": "modem",
+    "663": "monastery",
+    "664": "monitor",
+    "665": "moped",
+    "666": "mortar",
+    "667": "mortarboard",
+    "668": "mosque",
+    "669": "mosquito net",
+    "670": "motor scooter, scooter",
+    "671": "mountain bike, all-terrain bike, off-roader",
+    "672": "mountain tent",
+    "673": "mouse, computer mouse",
+    "674": "mousetrap",
+    "675": "moving van",
+    "676": "muzzle",
+    "677": "nail",
+    "678": "neck brace",
+    "679": "necklace",
+    "680": "nipple",
+    "681": "notebook, notebook computer",
+    "682": "obelisk",
+    "683": "oboe, hautboy, hautbois",
+    "684": "ocarina, sweet potato",
+    "685": "odometer, hodometer, mileometer, milometer",
+    "686": "oil filter",
+    "687": "organ, pipe organ",
+    "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO",
+    "689": "overskirt",
+    "690": "oxcart",
+    "691": "oxygen mask",
+    "692": "packet",
+    "693": "paddle, boat paddle",
+    "694": "paddlewheel, paddle wheel",
+    "695": "padlock",
+    "696": "paintbrush",
+    "697": "pajama, pyjama, pjs, jammies",
+    "698": "palace",
+    "699": "panpipe, pandean pipe, syrinx",
+    "700": "paper towel",
+    "701": "parachute, chute",
+    "702": "parallel bars, bars",
+    "703": "park bench",
+    "704": "parking meter",
+    "705": "passenger car, coach, carriage",
+    "706": "patio, terrace",
+    "707": "pay-phone, pay-station",
+    "708": "pedestal, plinth, footstall",
+    "709": "pencil box, pencil case",
+    "710": "pencil sharpener",
+    "711": "perfume, essence",
+    "712": "Petri dish",
+    "713": "photocopier",
+    "714": "pick, plectrum, plectron",
+    "715": "pickelhaube",
+    "716": "picket fence, paling",
+    "717": "pickup, pickup truck",
+    "718": "pier",
+    "719": "piggy bank, penny bank",
+    "720": "pill bottle",
+    "721": "pillow",
+    "722": "ping-pong ball",
+    "723": "pinwheel",
+    "724": "pirate, pirate ship",
+    "725": "pitcher, ewer",
+    "726": "plane, carpenters plane, woodworking plane",
+    "727": "planetarium",
+    "728": "plastic bag",
+    "729": "plate rack",
+    "730": "plow, plough",
+    "731": "plunger, plumbers helper",
+    "732": "Polaroid camera, Polaroid Land camera",
+    "733": "pole",
+    "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria",
+    "735": "poncho",
+    "736": "pool table, billiard table, snooker table",
+    "737": "pop bottle, soda bottle",
+    "738": "pot, flowerpot",
+    "739": "potters wheel",
+    "740": "power drill",
+    "741": "prayer rug, prayer mat",
+    "742": "printer",
+    "743": "prison, prison house",
+    "744": "projectile, missile",
+    "745": "projector",
+    "746": "puck, hockey puck",
+    "747": "punching bag, punch bag, punching ball, punchball",
+    "748": "purse",
+    "749": "quill, quill pen",
+    "750": "quilt, comforter, comfort, puff",
+    "751": "racer, race car, racing car",
+    "752": "racket, racquet",
+    "753": "radiator",
+    "754": "radio, wireless",
+    "755": "radio telescope, radio reflector",
+    "756": "rain barrel",
+    "757": "recreational vehicle, RV, R.V.",
+    "758": "reel",
+    "759": "reflex camera",
+    "760": "refrigerator, icebox",
+    "761": "remote control, remote",
+    "762": "restaurant, eating house, eating place, eatery",
+    "763": "revolver, six-gun, six-shooter",
+    "764": "rifle",
+    "765": "rocking chair, rocker",
+    "766": "rotisserie",
+    "767": "rubber eraser, rubber, pencil eraser",
+    "768": "rugby ball",
+    "769": "rule, ruler",
+    "770": "running shoe",
+    "771": "safe",
+    "772": "safety pin",
+    "773": "saltshaker, salt shaker",
+    "774": "sandal",
+    "775": "sarong",
+    "776": "sax, saxophone",
+    "777": "scabbard",
+    "778": "scale, weighing machine",
+    "779": "school bus",
+    "780": "schooner",
+    "781": "scoreboard",
+    "782": "screen, CRT screen",
+    "783": "screw",
+    "784": "screwdriver",
+    "785": "seat belt, seatbelt",
+    "786": "sewing machine",
+    "787": "shield, buckler",
+    "788": "shoe shop, shoe-shop, shoe store",
+    "789": "shoji",
+    "790": "shopping basket",
+    "791": "shopping cart",
+    "792": "shovel",
+    "793": "shower cap",
+    "794": "shower curtain",
+    "795": "ski",
+    "796": "ski mask",
+    "797": "sleeping bag",
+    "798": "slide rule, slipstick",
+    "799": "sliding door",
+    "800": "slot, one-armed bandit",
+    "801": "snorkel",
+    "802": "snowmobile",
+    "803": "snowplow, snowplough",
+    "804": "soap dispenser",
+    "805": "soccer ball",
+    "806": "sock",
+    "807": "solar dish, solar collector, solar furnace",
+    "808": "sombrero",
+    "809": "soup bowl",
+    "810": "space bar",
+    "811": "space heater",
+    "812": "space shuttle",
+    "813": "spatula",
+    "814": "speedboat",
+    "815": "spider web, spiders web",
+    "816": "spindle",
+    "817": "sports car, sport car",
+    "818": "spotlight, spot",
+    "819": "stage",
+    "820": "steam locomotive",
+    "821": "steel arch bridge",
+    "822": "steel drum",
+    "823": "stethoscope",
+    "824": "stole",
+    "825": "stone wall",
+    "826": "stopwatch, stop watch",
+    "827": "stove",
+    "828": "strainer",
+    "829": "streetcar, tram, tramcar, trolley, trolley car",
+    "830": "stretcher",
+    "831": "studio couch, day bed",
+    "832": "stupa, tope",
+    "833": "submarine, pigboat, sub, U-boat",
+    "834": "suit, suit of clothes",
+    "835": "sundial",
+    "836": "sunglass",
+    "837": "sunglasses, dark glasses, shades",
+    "838": "sunscreen, sunblock, sun blocker",
+    "839": "suspension bridge",
+    "840": "swab, swob, mop",
+    "841": "sweatshirt",
+    "842": "swimming trunks, bathing trunks",
+    "843": "swing",
+    "844": "switch, electric switch, electrical switch",
+    "845": "syringe",
+    "846": "table lamp",
+    "847": "tank, army tank, armored combat vehicle, armoured combat vehicle",
+    "848": "tape player",
+    "849": "teapot",
+    "850": "teddy, teddy bear",
+    "851": "television, television system",
+    "852": "tennis ball",
+    "853": "thatch, thatched roof",
+    "854": "theater curtain, theatre curtain",
+    "855": "thimble",
+    "856": "thresher, thrasher, threshing machine",
+    "857": "throne",
+    "858": "tile roof",
+    "859": "toaster",
+    "860": "tobacco shop, tobacconist shop, tobacconist",
+    "861": "toilet seat",
+    "862": "torch",
+    "863": "totem pole",
+    "864": "tow truck, tow car, wrecker",
+    "865": "toyshop",
+    "866": "tractor",
+    "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi",
+    "868": "tray",
+    "869": "trench coat",
+    "870": "tricycle, trike, velocipede",
+    "871": "trimaran",
+    "872": "tripod",
+    "873": "triumphal arch",
+    "874": "trolleybus, trolley coach, trackless trolley",
+    "875": "trombone",
+    "876": "tub, vat",
+    "877": "turnstile",
+    "878": "typewriter keyboard",
+    "879": "umbrella",
+    "880": "unicycle, monocycle",
+    "881": "upright, upright piano",
+    "882": "vacuum, vacuum cleaner",
+    "883": "vase",
+    "884": "vault",
+    "885": "velvet",
+    "886": "vending machine",
+    "887": "vestment",
+    "888": "viaduct",
+    "889": "violin, fiddle",
+    "890": "volleyball",
+    "891": "waffle iron",
+    "892": "wall clock",
+    "893": "wallet, billfold, notecase, pocketbook",
+    "894": "wardrobe, closet, press",
+    "895": "warplane, military plane",
+    "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin",
+    "897": "washer, automatic washer, washing machine",
+    "898": "water bottle",
+    "899": "water jug",
+    "900": "water tower",
+    "901": "whiskey jug",
+    "902": "whistle",
+    "903": "wig",
+    "904": "window screen",
+    "905": "window shade",
+    "906": "Windsor tie",
+    "907": "wine bottle",
+    "908": "wing",
+    "909": "wok",
+    "910": "wooden spoon",
+    "911": "wool, woolen, woollen",
+    "912": "worm fence, snake fence, snake-rail fence, Virginia fence",
+    "913": "wreck",
+    "914": "yawl",
+    "915": "yurt",
+    "916": "web site, website, internet site, site",
+    "917": "comic book",
+    "918": "crossword puzzle, crossword",
+    "919": "street sign",
+    "920": "traffic light, traffic signal, stoplight",
+    "921": "book jacket, dust cover, dust jacket, dust wrapper",
+    "922": "menu",
+    "923": "plate",
+    "924": "guacamole",
+    "925": "consomme",
+    "926": "hot pot, hotpot",
+    "927": "trifle",
+    "928": "ice cream, icecream",
+    "929": "ice lolly, lolly, lollipop, popsicle",
+    "930": "French loaf",
+    "931": "bagel, beigel",
+    "932": "pretzel",
+    "933": "cheeseburger",
+    "934": "hotdog, hot dog, red hot",
+    "935": "mashed potato",
+    "936": "head cabbage",
+    "937": "broccoli",
+    "938": "cauliflower",
+    "939": "zucchini, courgette",
+    "940": "spaghetti squash",
+    "941": "acorn squash",
+    "942": "butternut squash",
+    "943": "cucumber, cuke",
+    "944": "artichoke, globe artichoke",
+    "945": "bell pepper",
+    "946": "cardoon",
+    "947": "mushroom",
+    "948": "Granny Smith",
+    "949": "strawberry",
+    "950": "orange",
+    "951": "lemon",
+    "952": "fig",
+    "953": "pineapple, ananas",
+    "954": "banana",
+    "955": "jackfruit, jak, jack",
+    "956": "custard apple",
+    "957": "pomegranate",
+    "958": "hay",
+    "959": "carbonara",
+    "960": "chocolate sauce, chocolate syrup",
+    "961": "dough",
+    "962": "meat loaf, meatloaf",
+    "963": "pizza, pizza pie",
+    "964": "potpie",
+    "965": "burrito",
+    "966": "red wine",
+    "967": "espresso",
+    "968": "cup",
+    "969": "eggnog",
+    "970": "alp",
+    "971": "bubble",
+    "972": "cliff, drop, drop-off",
+    "973": "coral reef",
+    "974": "geyser",
+    "975": "lakeside, lakeshore",
+    "976": "promontory, headland, head, foreland",
+    "977": "sandbar, sand bar",
+    "978": "seashore, coast, seacoast, sea-coast",
+    "979": "valley, vale",
+    "980": "volcano",
+    "981": "ballplayer, baseball player",
+    "982": "groom, bridegroom",
+    "983": "scuba diver",
+    "984": "rapeseed",
+    "985": "daisy",
+    "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",
+    "987": "corn",
+    "988": "acorn",
+    "989": "hip, rose hip, rosehip",
+    "990": "buckeye, horse chestnut, conker",
+    "991": "coral fungus",
+    "992": "agaric",
+    "993": "gyromitra",
+    "994": "stinkhorn, carrion fungus",
+    "995": "earthstar",
+    "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa",
+    "997": "bolete",
+    "998": "ear, spike, capitulum",
+    "999": "toilet tissue, toilet paper, bathroom tissue"
+  }
+}

W-Flow-L/pipeline.py

@@ -0,0 +1,754 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.util
+import inspect
+import json
+import sys
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import deprecate, is_torch_xla_available, logging, replace_example_docstring
+from PIL import Image
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from diffusers import DiffusionPipeline
+
+        >>> model_dir = Path("W-Flow-diffusers/W-Flow-L").resolve()
+        >>> pipe = DiffusionPipeline.from_pretrained(
+        ...     str(model_dir),
+        ...     local_files_only=True,
+        ...     custom_pipeline=str(model_dir / "pipeline.py"),
+        ...     trust_remote_code=True,
+        ...     torch_dtype=torch.bfloat16,
+        ... ).to("cuda")
+
+        >>> image = pipe(
+        ...     class_labels="golden retriever",
+        ...     num_inference_steps=1,
+        ...     guidance_scale=1.14,
+        ...     generator=torch.Generator(device="cuda").manual_seed(1),
+        ... ).images[0]
+        >>> image.save("demo.png")
+
+        >>> # Hugging Face Hub style model id: UserID/RepoID
+        >>> # pipe = DiffusionPipeline.from_pretrained(
+        >>> #     "BiliSakura/W-Flow-diffusers",
+        >>> #     subfolder="W-Flow-L",
+        >>> #     custom_pipeline="pipeline.py",
+        >>> #     trust_remote_code=True,
+        >>> #     torch_dtype=torch.bfloat16,
+        >>> # )
+        ```
+"""
+
+
+class WFlowPipeline(DiffusionPipeline):
+    r"""
+    Class-conditional one-step pipeline for W-Flow.
+
+    W-Flow maps SD-VAE latent noise to a clean 32×32 latent in one forward pass, then decodes it
+    with `stabilityai/sd-vae-ft-mse` to a 256×256 image. Guidance is embedded in the generator
+    (the official `cfg_scale`, which is `w + 1` in the paper). There is no unconditional class
+    and no second forward.
+
+    The pipeline inherits from [`DiffusionPipeline`]. Check the superclass documentation for the
+    generic methods (download, save, device placement, CPU offload).
+
+    Args:
+        transformer ([`ModelMixin`]):
+            [`WFlowTransformer2DModel`] that predicts the clean latent.
+        vae ([`AutoencoderKL`]):
+            SD-VAE decoder. Released checkpoints use `stabilityai/sd-vae-ft-mse` and scaling 0.18215.
+        scheduler ([`WFlowScheduler`]):
+            One-step scheduler. Loaded from `scheduler/`.
+        id2label (`dict[int, str]`, *optional*):
+            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
+            Loaded from `model_index.json` when omitted.
+
+        The default guidance scale lives on `transformer.config.guidance_scale`.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+    _callback_tensor_inputs = ["latents"]
+
+    def __init__(
+        self,
+        transformer: ModelMixin,
+        vae: AutoencoderKL,
+        scheduler: SchedulerMixin,
+        id2label: dict[int, str] | dict[str, str] | None = None,
+    ):
+        super().__init__()
+        if scheduler is None:
+            scheduler = self._default_scheduler()
+        self._reject_karras_scheduler(scheduler)
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+        normalized_labels = self._normalize_id2label(id2label)
+        self.register_to_config(id2label={str(key): value for key, value in normalized_labels.items()})
+        self._id2label = normalized_labels
+        self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = bool(self._id2label)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    def _sampling_default(self, name: str, fallback: float) -> float:
+        r"""
+        Read a sampling default from the transformer config.
+
+        Args:
+            name (`str`):
+                Key on `transformer.config`, such as `guidance_scale`.
+            fallback (`float`):
+                Value used when the transformer config does not define `name`.
+
+        Returns:
+            `float`: The configured value, or `fallback`.
+        """
+        value = getattr(self.transformer.config, name, None)
+        if value is None:
+            return fallback
+        return value
+
+    @staticmethod
+    def _reject_karras_scheduler(scheduler: SchedulerMixin) -> None:
+        r"""
+        Reject multi-step Karras schedulers.
+
+        W-Flow predicts the clean latent in one evaluation and has no diffusion timestep, so a
+        Karras noise schedule cannot be integrated.
+
+        Args:
+            scheduler ([`SchedulerMixin`]):
+                Scheduler attached to the pipeline.
+        """
+        scheduler_name = scheduler.__class__.__name__
+        if scheduler_name in KarrasDiffusionSchedulers.__members__:
+            raise TypeError(
+                "W-Flow is a one-step generator and does not integrate a Karras noise schedule. "
+                f"Received {scheduler_name}. Use WFlowScheduler."
+            )
+
+    @staticmethod
+    def _default_scheduler() -> SchedulerMixin:
+        r"""
+        Build [`WFlowScheduler`] from the variant directory.
+
+        Returns:
+            [`WFlowScheduler`]: One-step sample scheduler.
+        """
+        module_path = Path(__file__).resolve().parent / "scheduler" / "scheduling_wflow.py"
+        spec = importlib.util.spec_from_file_location("scheduling_wflow", module_path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Could not load the W-Flow scheduler from {module_path}.")
+        module = importlib.util.module_from_spec(spec)
+        sys.modules["scheduling_wflow"] = module
+        spec.loader.exec_module(module)
+        return module.WFlowScheduler()
+
+    def _ensure_labels_loaded(self) -> None:
+        r"""
+        Fill `id2label` from `model_index.json` when the constructor did not receive it.
+        """
+        if self._labels_loaded_from_model_index:
+            return
+        loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None))
+        if loaded:
+            self._id2label = loaded
+            self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = True
+
+    @staticmethod
+    def _normalize_id2label(id2label: dict[int, str] | dict[str, str] | None) -> dict[int, str]:
+        r"""
+        Cast label-map keys to integers.
+
+        Args:
+            id2label (`dict`, *optional*):
+                Class id to English label. Keys may be strings, which is how JSON stores them.
+
+        Returns:
+            `dict[int, str]`: Integer-keyed map. Empty when `id2label` is missing.
+        """
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _read_id2label_from_model_index(variant_path: str | None) -> dict[int, str]:
+        r"""
+        Read `id2label` from a local variant directory.
+
+        Args:
+            variant_path (`str`, *optional*):
+                Directory that contains `model_index.json`.
+
+        Returns:
+            `dict[int, str]`: Loaded map, or an empty dict when the file or key is absent.
+        """
+        if not variant_path:
+            return {}
+        model_index_path = Path(variant_path).resolve() / "model_index.json"
+        if not model_index_path.exists():
+            return {}
+        raw = json.loads(model_index_path.read_text(encoding="utf-8"))
+        id2label = raw.get("id2label")
+        if not isinstance(id2label, dict):
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _build_label2id(id2label: dict[int, str]) -> dict[str, int]:
+        r"""
+        Invert `id2label`, splitting comma-separated synonyms.
+
+        Args:
+            id2label (`dict[int, str]`):
+                Class id to one or more English names.
+
+        Returns:
+            `dict[str, int]`: Synonym to class id, sorted by synonym.
+        """
+        label2id: dict[str, int] = {}
+        for class_id, value in id2label.items():
+            for synonym in value.split(","):
+                synonym = synonym.strip()
+                if synonym:
+                    label2id[synonym] = int(class_id)
+        return dict(sorted(label2id.items()))
+
+    @property
+    def interrupt(self) -> bool:
+        r"""
+        Whether the current call was asked to stop.
+
+        Returns:
+            `bool`: The `_interrupt` flag. `__call__` sets it to `False` at the start. Assign
+            `_interrupt = True` during a call to skip the remaining forward.
+        """
+        return self._interrupt
+
+    @property
+    def id2label(self) -> dict[int, str]:
+        r"""
+        ImageNet id to English label map.
+
+        Returns:
+            `dict[int, str]`: The map loaded from the constructor or `model_index.json`.
+        """
+        self._ensure_labels_loaded()
+        return self._id2label
+
+    def get_label_ids(self, label: str | list[str]) -> list[int]:
+        r"""
+        Map ImageNet English labels to class ids.
+
+        Args:
+            label (`str` or `list[str]`):
+                One label or a list of labels. Each string must match a synonym in `id2label`
+                after stripping whitespace. Synonyms are the comma-separated pieces of each value.
+
+        Returns:
+            `list[int]`: Class ids in the same order as `label`.
+        """
+        self._ensure_labels_loaded()
+        label2id = self.labels
+        if not label2id:
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
+        if isinstance(label, str):
+            label = [label]
+        missing = [item for item in label if item not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
+        return [label2id[item] for item in label]
+
+    def _normalize_class_labels(self, class_labels: int | str | list[int | str]) -> list[int]:
+        r"""
+        Turn a class-label argument into a list of integer ids.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                One id, one English label, or a list of either. A list must not mix types.
+
+        Returns:
+            `list[int]`: One id per requested image.
+        """
+        if isinstance(class_labels, int):
+            return [class_labels]
+        if isinstance(class_labels, str):
+            return self.get_label_ids(class_labels)
+        if len(class_labels) == 0:
+            raise ValueError("`class_labels` must contain at least one class id or label.")
+        if isinstance(class_labels[0], str):
+            if not all(isinstance(item, str) for item in class_labels):
+                raise TypeError("`class_labels` must be all strings or all integers.")
+            return self.get_label_ids(class_labels)
+        if not all(isinstance(item, int) for item in class_labels):
+            raise TypeError("`class_labels` must be all strings or all integers.")
+        return list(class_labels)
+
+    def prepare_extra_step_kwargs(
+        self,
+        generator: torch.Generator | list[torch.Generator] | None,
+        eta: float,
+    ) -> dict[str, Any]:
+        r"""
+        Extra arguments forwarded to `scheduler.step`, depending on what that method accepts.
+
+        Args:
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Generator passed through when the scheduler step samples noise.
+            eta (`float`):
+                DDIM eta in `[0, 1]`. Ignored by [`WFlowScheduler`], whose `step` has no `eta` argument.
+
+        Returns:
+            `dict`: Keyword arguments for `scheduler.step`.
+        """
+        extra_step_kwargs: dict[str, Any] = {}
+        step_params = set(inspect.signature(self.scheduler.step).parameters.keys())
+        if "eta" in step_params:
+            extra_step_kwargs["eta"] = eta
+        if "generator" in step_params:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        class_labels: int | str | list[int | str],
+        height: int,
+        width: int,
+        num_inference_steps: int,
+        guidance_scale: float,
+        noise_scale: float,
+        output_type: str,
+        callback_steps: int | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+    ) -> None:
+        r"""
+        Reject arguments that the generator cannot run.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                Class ids or English labels. A list must be non-empty.
+            height (`int`):
+                Output height in pixels. Must be the pretrained resolution, 256.
+            width (`int`):
+                Output width in pixels. Must be the pretrained resolution, 256.
+            num_inference_steps (`int`):
+                Must be 1. W-Flow evaluates the network once.
+            guidance_scale (`float`):
+                Embedded guidance (`w + 1` in the paper). Must be finite and positive.
+            noise_scale (`float`):
+                Scale of the reference Gaussian, the official `temp`. Must be positive.
+            output_type (`str`):
+                One of `"pil"`, `"np"`, `"pt"`, `"latent"`.
+            callback_steps (`int`, *optional*):
+                Deprecated callback period. When set, must be a positive integer.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names the step callback may read. Each name must be in `_callback_tensor_inputs`.
+        """
+        if isinstance(class_labels, list) and len(class_labels) == 0:
+            raise ValueError("`class_labels` must contain at least one class id or label.")
+        native = int(self.transformer.config.sample_size) * int(self.vae_scale_factor)
+        if height != native or width != native:
+            raise ValueError(
+                f"`height` and `width` must be the pretrained resolution {native}. "
+                f"The patch embedding is fixed to a {self.transformer.config.sample_size} latent. Got {(height, width)}."
+            )
+        if num_inference_steps != 1:
+            raise ValueError(
+                "W-Flow generates in one forward pass. "
+                f"`num_inference_steps` must be 1, got {num_inference_steps}."
+            )
+        if not torch.isfinite(torch.tensor(float(guidance_scale))) or float(guidance_scale) <= 0:
+            raise ValueError(f"`guidance_scale` must be finite and positive, got {guidance_scale}.")
+        if noise_scale <= 0:
+            raise ValueError(f"`noise_scale` must be positive, got {noise_scale}.")
+        if output_type not in {"pil", "np", "pt", "latent"}:
+            raise ValueError(f"Unsupported `output_type` {output_type!r}. Choose from 'pil', 'np', 'pt', 'latent'.")
+        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps < 1):
+            raise ValueError(f"`callback_steps` must be a positive integer, got {callback_steps}.")
+        if callback_on_step_end_tensor_inputs is not None:
+            unknown = [name for name in callback_on_step_end_tensor_inputs if name not in self._callback_tensor_inputs]
+            if unknown:
+                raise ValueError(
+                    f"`callback_on_step_end_tensor_inputs` contains {unknown}, which are not in "
+                    f"{self._callback_tensor_inputs}."
+                )
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels: int,
+        height: int,
+        width: int,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None,
+        noise_labels: torch.Tensor | None,
+        noise_scale: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Draw reference noise in the official generator order.
+
+        The official sampler draws `(batch, height, width, channels)` first, then integer noise
+        codes, from the same generator. This method follows that order so a seeded generator
+        matches `DitGen.forward`.
+
+        Args:
+            batch_size (`int`):
+                Number of images.
+            num_channels (`int`):
+                Latent channels. Released checkpoints use 4.
+            height (`int`):
+                Latent height. Released checkpoints use 32.
+            width (`int`):
+                Latent width. Released checkpoints use 32.
+            device (`torch.device`):
+                Device of the returned tensors.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG. A list must have one generator per image. Each generator is consumed in the
+                order Gaussian noise, then noise codes.
+            latents (`torch.Tensor`, *optional*):
+                Reference noise `(batch, channels, height, width)`. When set, it is not rescaled.
+            noise_labels (`torch.Tensor`, *optional*):
+                Integer codes `(batch, noise_coords)`. When omitted, they are drawn here.
+            noise_scale (`float`):
+                Multiplier on the Gaussian, the official `temp`. Ignored when `latents` is passed.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`:
+                Reference noise in BCHW float32, and noise codes `(batch, noise_coords)`.
+        """
+        noise_classes = int(self.transformer.config.noise_classes)
+        noise_coords = int(self.transformer.config.noise_coords)
+        if latents is None:
+            if isinstance(generator, list):
+                if len(generator) != batch_size:
+                    raise ValueError(f"Got {len(generator)} generators for a batch of {batch_size}.")
+                pieces = [
+                    torch.randn((1, height, width, num_channels), generator=one, device=device) for one in generator
+                ]
+                drawn = torch.cat(pieces, dim=0)
+            else:
+                drawn = torch.randn(
+                    (batch_size, height, width, num_channels),
+                    generator=generator,
+                    device=device,
+                )
+            latents = drawn.permute(0, 3, 1, 2).contiguous() * float(noise_scale)
+        elif latents.shape != (batch_size, num_channels, height, width):
+            raise ValueError(
+                f"`latents` must have shape {(batch_size, num_channels, height, width)}, got {tuple(latents.shape)}."
+            )
+        else:
+            latents = latents.to(device=device)
+
+        if noise_labels is None:
+            if isinstance(generator, list) and latents is not None:
+                # When latents were supplied, generators have not been consumed yet only if the
+                # caller passed both. Drawing labels still uses those generators.
+                codes = [
+                    torch.randint(0, noise_classes, (1, noise_coords), generator=one, device=device)
+                    for one in generator
+                ]
+                noise_labels = torch.cat(codes, dim=0)
+            else:
+                noise_labels = torch.randint(
+                    0,
+                    max(1, noise_classes),
+                    (batch_size, max(1, noise_coords)),
+                    generator=None if isinstance(generator, list) else generator,
+                    device=device,
+                )
+        noise_labels = noise_labels.to(device=device, dtype=torch.long)
+        if noise_labels.shape != (batch_size, noise_coords):
+            raise ValueError(
+                f"`noise_labels` must have shape {(batch_size, noise_coords)}, got {tuple(noise_labels.shape)}."
+            )
+        return latents.float(), noise_labels
+
+    def _prepare_scheduler(self, num_inference_steps: int, device: torch.device) -> torch.Tensor:
+        r"""
+        Install the one-step schedule and return its timesteps.
+
+        Args:
+            num_inference_steps (`int`):
+                Must be 1.
+            device (`torch.device`):
+                Device of the timestep tensor.
+
+        Returns:
+            `torch.Tensor`: The single timestep.
+        """
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        return self.scheduler.timesteps
+
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        r"""
+        Decode scaled SD-VAE latents to pixels in `[-1, 1]`.
+
+        Args:
+            latents (`torch.Tensor`):
+                Clean latents `(batch, 4, 32, 32)` in the training scale (`* 0.18215`).
+
+        Returns:
+            `torch.Tensor`: Images `(batch, 3, height, width)` in `[-1, 1]`.
+        """
+        scaling = float(getattr(self.vae.config, "scaling_factor", 0.18215))
+        needs_upcast = bool(getattr(self.vae.config, "force_upcast", False)) and self.vae.dtype != torch.float32
+        original_dtype = self.vae.dtype
+        if needs_upcast:
+            self.vae.to(dtype=torch.float32)
+        decoded = self.vae.decode((latents / scaling).to(dtype=self.vae.dtype), return_dict=False)[0]
+        if needs_upcast:
+            self.vae.to(dtype=original_dtype)
+        return decoded
+
+    def _images_from_pixels(self, pixels: torch.Tensor, output_type: str) -> list[Image.Image] | torch.Tensor | list:
+        r"""
+        Convert decoded pixels to the requested output type.
+
+        Args:
+            pixels (`torch.Tensor`):
+                Images `(batch, 3, height, width)` in `[-1, 1]`.
+            output_type (`str`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"`.
+
+        Returns:
+            Images in the requested type.
+        """
+        images = ((pixels.float() + 1) / 2).clamp(0, 1)
+        if output_type == "pt":
+            return images
+        images = (images * 255).round().to(torch.uint8)
+        arrays = images.permute(0, 2, 3, 1).cpu().numpy()
+        if output_type == "np":
+            return arrays
+        return [Image.fromarray(image) for image in arrays]
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        class_labels: int | str | list[int | str],
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        guidance_scale: float | None = None,
+        noise_scale: float = 1.0,
+        eta: float = 0.0,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        noise_labels: torch.Tensor | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Callable[[int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+        **kwargs,
+    ) -> ImagePipelineOutput | tuple:
+        r"""
+        Generate images from ImageNet class labels.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                One ImageNet class id, one English label, or a list of either. English labels are
+                resolved with [`get_label_ids`]. The list length is the batch size.
+            height (`int`, *optional*):
+                Image height in pixels. Defaults to 256 (`sample_size * vae_scale_factor`). Other
+                sizes are rejected because the patch embedding is fixed.
+            width (`int`, *optional*):
+                Image width in pixels. Defaults to 256.
+            num_inference_steps (`int`, *optional*):
+                Must be 1. Defaults to `transformer.config.default_num_inference_steps`.
+            guidance_scale (`float`, *optional*):
+                Official CFG value, equal to `w + 1` in the paper. Defaults to
+                `transformer.config.guidance_scale` (1.19 / 1.14 / 1.09 / 1.5 by variant). This
+                scale is embedded; the pipeline does not run an unconditional branch.
+            noise_scale (`float`, defaults to 1.0):
+                Multiplier on the reference Gaussian. The official sampler calls this `temp`.
+                Ignored when `latents` is passed.
+            eta (`float`, defaults to 0.0):
+                DDIM eta. Ignored by [`WFlowScheduler`].
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG for the reference noise and the discrete noise codes. One generator is consumed
+                in that order, matching the official sampler.
+            latents (`torch.Tensor`, *optional*):
+                Reference noise `(batch, 4, 32, 32)`. Not multiplied by `noise_scale`.
+            noise_labels (`torch.Tensor`, *optional*):
+                Discrete noise codes `(batch, noise_coords)`. Drawn from `generator` when omitted.
+            output_type (`str`, defaults to `"pil"`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"` (scaled SD-VAE latent).
+            return_dict (`bool`, defaults to `True`):
+                Return [`ImagePipelineOutput`] when `True`, otherwise a one-tuple of images.
+            callback_on_step_end (`Callable` or `PipelineCallback`, *optional*):
+                Called as `callback_on_step_end(pipeline, step, timestep, callback_kwargs)` after the
+                scheduler step. Return a dict to replace tensors listed in
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names passed to the step callback. Defaults to `["latents"]`.
+
+        Examples:
+
+        Returns:
+            [`ImagePipelineOutput`] or `tuple`:
+                When `return_dict` is `True`, [`ImagePipelineOutput`] with the images. Otherwise a tuple
+                whose first element is the images.
+        """
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        if kwargs:
+            raise TypeError(f"Unexpected arguments: {sorted(kwargs)}.")
+        if callback_on_step_end_tensor_inputs is None:
+            callback_on_step_end_tensor_inputs = ["latents"]
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        image_size = int(self.transformer.config.sample_size) * int(self.vae_scale_factor)
+        height = image_size if height is None else int(height)
+        width = image_size if width is None else int(width)
+        if num_inference_steps is None:
+            num_inference_steps = int(self._sampling_default("default_num_inference_steps", 1))
+        if guidance_scale is None:
+            guidance_scale = float(self._sampling_default("guidance_scale", 1.0))
+
+        # 1. Check inputs.
+        self.check_inputs(
+            class_labels,
+            height,
+            width,
+            num_inference_steps,
+            guidance_scale,
+            noise_scale,
+            output_type,
+            callback_steps,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = float(guidance_scale)
+        self._interrupt = False
+
+        # 2. Define call parameters.
+        class_label_ids = self._normalize_class_labels(class_labels)
+        batch_size = len(class_label_ids)
+        device = self._execution_device
+        num_classes = int(self.transformer.config.num_classes)
+        if any(class_id < 0 or class_id >= num_classes for class_id in class_label_ids):
+            raise ValueError(f"Class ids must lie in [0, {num_classes}), got {class_label_ids}.")
+
+        # 3. Encode input condition (class ids; guidance is embedded inside the transformer).
+        class_labels_t = torch.tensor(class_label_ids, device=device, dtype=torch.long)
+
+        # 4. Prepare timesteps.
+        scheduled_timesteps = self._prepare_scheduler(num_inference_steps, device)
+
+        # 5. Prepare latent variables. Noise codes are drawn here so one generator matches DitGen.
+        latent_height = height // int(self.vae_scale_factor)
+        latent_width = width // int(self.vae_scale_factor)
+        latents, noise_labels = self.prepare_latents(
+            batch_size,
+            int(self.transformer.config.in_channels),
+            latent_height,
+            latent_width,
+            device,
+            generator,
+            latents,
+            noise_labels,
+            noise_scale,
+        )
+
+        # 6. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_warmup_steps = len(scheduled_timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(scheduled_timesteps)
+
+        # 7. Denoising loop.
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, timestep in enumerate(scheduled_timesteps):
+                if self.interrupt:
+                    continue
+
+                latent_model_input = self.scheduler.scale_model_input(latents, timestep)
+                model_output = self.transformer(
+                    latent_model_input,
+                    class_labels=class_labels_t,
+                    guidance_scale=self._guidance_scale,
+                    noise_labels=noise_labels,
+                    return_dict=False,
+                )[0]
+                latents = self.scheduler.step(
+                    model_output,
+                    timestep,
+                    latents,
+                    **extra_step_kwargs,
+                    return_dict=False,
+                )[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {key: locals()[key] for key in callback_on_step_end_tensor_inputs}
+                    callback_outputs = callback_on_step_end(self, i, timestep, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+
+                if i == len(scheduled_timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, timestep, latents)
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            images = latents
+        else:
+            images = self._images_from_pixels(self._decode_latents(latents), output_type)
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

W-Flow-L/scheduler/scheduler_config.json

@@ -0,0 +1,6 @@
+{
+  "_class_name": "WFlowScheduler",
+  "_diffusers_version": "0.39.0",
+  "num_train_timesteps": 1000,
+  "prediction_type": "sample"
+}

W-Flow-L/scheduler/scheduling_wflow.py

@@ -0,0 +1,152 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class WFlowSchedulerOutput(BaseOutput):
+    r"""
+    Output of [`WFlowScheduler.step`].
+
+    Args:
+        prev_sample (`torch.Tensor`):
+            The model prediction. W-Flow returns the clean latent in one evaluation.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class WFlowScheduler(SchedulerMixin, ConfigMixin):
+    r"""
+    One-step scheduler for W-Flow.
+
+    The transformer predicts the clean latent directly. `set_timesteps` accepts a single evaluation,
+    and `step` returns that prediction. `scale_model_input` leaves the reference noise unchanged.
+    A Karras noise schedule does not apply: the network has no diffusion timestep.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            Stored for config compatibility. W-Flow does not index it.
+        prediction_type (`str`, defaults to `"sample"`):
+            The model output is the clean sample.
+    """
+
+    order = 1
+    _compatibles: list[str] = []
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        prediction_type: str = "sample",
+    ):
+        if num_train_timesteps < 1:
+            raise ValueError(f"num_train_timesteps must be >= 1, got {num_train_timesteps}.")
+        if prediction_type != "sample":
+            raise ValueError(f"WFlowScheduler prediction_type must be 'sample', got {prediction_type!r}.")
+        self.timesteps: torch.Tensor | None = None
+        self.num_inference_steps: int | None = None
+        self._step_index: int | None = None
+
+    @property
+    def init_noise_sigma(self) -> float:
+        r"""
+        Scale of the standard normal used to start sampling.
+
+        Returns:
+            `float`: `1.0`. The pipeline multiplies this by `noise_scale`.
+        """
+        return 1.0
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
+        r"""
+        Return the sample unchanged.
+
+        Args:
+            sample (`torch.Tensor`):
+                Reference noise `(batch, channels, height, width)`.
+            timestep (`float` or `torch.Tensor`):
+                Unused. W-Flow has no diffusion timestep.
+
+        Returns:
+            `torch.Tensor`: `sample`.
+        """
+        del timestep
+        return sample
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: str | torch.device | None = None,
+    ) -> None:
+        r"""
+        Install the single evaluation step.
+
+        Args:
+            num_inference_steps (`int`):
+                Must be 1. W-Flow generates the latent in one forward pass.
+            device (`str` or `torch.device`, *optional*):
+                Device of `self.timesteps`.
+        """
+        if num_inference_steps != 1:
+            raise ValueError(
+                "W-Flow generates in one forward pass. "
+                f"num_inference_steps must be 1, got {num_inference_steps}."
+            )
+        self.num_inference_steps = 1
+        self.timesteps = torch.tensor([0.0], device=device, dtype=torch.float32)
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: float | torch.Tensor,
+        sample: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        return_dict: bool = True,
+    ) -> WFlowSchedulerOutput | tuple[torch.Tensor]:
+        r"""
+        Replace the reference noise with the predicted clean latent.
+
+        Args:
+            model_output (`torch.Tensor`):
+                Clean-latent prediction `(batch, channels, height, width)`.
+            timestep (`float` or `torch.Tensor`):
+                The single scheduled timestep. Accepted for the Diffusers step signature.
+            sample (`torch.Tensor`):
+                Current reference noise. Same shape as `model_output`.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Accepted so pipelines can pass it through. W-Flow does not draw noise in `step`.
+            return_dict (`bool`, defaults to `True`):
+                Return [`WFlowSchedulerOutput`] when `True`, otherwise a one-tuple.
+
+        Returns:
+            [`WFlowSchedulerOutput`] or `tuple`:
+                `prev_sample` is `model_output`.
+        """
+        del timestep, sample, generator
+        if self.timesteps is None:
+            raise ValueError("Call set_timesteps before step.")
+        if self._step_index is None or self._step_index > 0:
+            raise ValueError("WFlowScheduler has already consumed its single step. Call set_timesteps again.")
+        self._step_index = 1
+        if not return_dict:
+            return (model_output,)
+        return WFlowSchedulerOutput(prev_sample=model_output)

W-Flow-L/transformer/config.json

@@ -0,0 +1,26 @@
+{
+  "_class_name": "WFlowTransformer2DModel",
+  "_diffusers_version": "0.39.0",
+  "attn_fp32": false,
+  "cond_dim": 1024,
+  "default_num_inference_steps": 1,
+  "depth": 24,
+  "guidance_scale": 1.14,
+  "hidden_size": 1024,
+  "in_channels": 4,
+  "mlp_ratio": 4.0,
+  "model_type": "W-Flow-L",
+  "n_cls_tokens": 16,
+  "noise_classes": 64,
+  "noise_coords": 32,
+  "num_attention_heads": 16,
+  "num_classes": 1000,
+  "out_channels": 4,
+  "patch_size": 2,
+  "sample_size": 32,
+  "use_bf16": true,
+  "use_qknorm": true,
+  "use_rmsnorm": true,
+  "use_rope": true,
+  "use_swiglu": true
+}

W-Flow-L/transformer/diffusion_pytorch_model.safetensors

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

W-Flow-L/transformer/transformer_wflow.py

@@ -0,0 +1,720 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class WFlowTransformer2DModelOutput(BaseOutput):
+    r"""
+    Output of [`WFlowTransformer2DModel`].
+
+    Args:
+        sample (`torch.Tensor`):
+            Predicted SD-VAE latent of shape `(batch, channels, height, width)`. Values are in the
+            scaled latent space used at training time (`encoder_mean * 0.18215`).
+    """
+
+    sample: torch.Tensor
+
+
+def _init_linear(linear: nn.Linear, weight_init: str = "xavier_uniform", bias_init: str = "zeros") -> None:
+    r"""
+    Initialize a linear layer the way the official generator does.
+
+    Args:
+        linear (`torch.nn.Linear`):
+            Layer to initialize in place.
+        weight_init (`str`, defaults to `"xavier_uniform"`):
+            `"xavier_uniform"`, `"zeros"`, or `"normal"` (std 0.02).
+        bias_init (`str`, defaults to `"zeros"`):
+            Bias fill. Only `"zeros"` is used by the released architecture.
+    """
+    if weight_init == "xavier_uniform":
+        nn.init.xavier_uniform_(linear.weight)
+    elif weight_init == "zeros":
+        nn.init.zeros_(linear.weight)
+    elif weight_init == "normal":
+        nn.init.normal_(linear.weight, std=0.02)
+    else:
+        raise ValueError(f"weight_init must be xavier_uniform, zeros, or normal, got {weight_init!r}.")
+    if linear.bias is not None:
+        if bias_init != "zeros":
+            raise ValueError(f"bias_init must be 'zeros', got {bias_init!r}.")
+        nn.init.zeros_(linear.bias)
+
+
+class _TorchLinear(nn.Module):
+    r"""
+    Linear wrapper whose parameter names match the official `TorchLinear` (`*.linear.weight`).
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        weight_init: str = "xavier_uniform",
+        bias_init: str = "zeros",
+    ):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        _init_linear(self.linear, weight_init=weight_init, bias_init=bias_init)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Apply the wrapped linear layer.
+
+        Args:
+            x (`torch.Tensor`):
+                Input of shape `(..., in_features)`.
+
+        Returns:
+            `torch.Tensor`: Output of shape `(..., out_features)`.
+        """
+        return self.linear(x)
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(int(dim)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Root-mean-square normalize the last dimension.
+
+        Args:
+            x (`torch.Tensor`):
+                Input of shape `(..., dim)`.
+
+        Returns:
+            `torch.Tensor`: Normalized tensor, same shape and dtype as `x`.
+        """
+        input_dtype = x.dtype
+        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
+        normalized = x * torch.rsqrt(variance + self.eps)
+        return (normalized * self.weight).to(dtype=input_dtype)
+
+
+def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    r"""
+    Apply an adaLN-Zero shift and scale.
+
+    Args:
+        x (`torch.Tensor`):
+            Token tensor `(batch, tokens, hidden)`.
+        shift (`torch.Tensor`):
+            Per-sample shift `(batch, hidden)`.
+        scale (`torch.Tensor`):
+            Per-sample scale `(batch, hidden)`.
+
+    Returns:
+        `torch.Tensor`: `x * (1 + scale) + shift`, broadcast over tokens.
+    """
+    return x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def _sincos_1d(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    r"""
+    Sine-cosine embedding of 1-D positions.
+
+    Args:
+        embed_dim (`int`):
+            Even embedding width.
+        pos (`numpy.ndarray`):
+            Positions, flattened inside the function.
+
+    Returns:
+        `numpy.ndarray`: Array of shape `(positions, embed_dim)`.
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError(f"embed_dim must be even, got {embed_dim}.")
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega = 1.0 / 10000 ** (omega / (embed_dim / 2.0))
+    out = np.einsum("m,d->md", pos.reshape(-1), omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+def _sincos_2d(embed_dim: int, grid_size: int) -> np.ndarray:
+    r"""
+    Sine-cosine embedding of a square grid, matching the official generator.
+
+    Args:
+        embed_dim (`int`):
+            Full embedding width. Half is used for each axis.
+        grid_size (`int`):
+            Tokens along one side of the latent grid.
+
+    Returns:
+        `numpy.ndarray`: Array of shape `(grid_size * grid_size, embed_dim)`.
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0).reshape(2, 1, grid_size, grid_size)
+    half = embed_dim // 2
+    return np.concatenate([_sincos_1d(half, grid[0]), _sincos_1d(half, grid[1])], axis=1)
+
+
+def _swiglu_hidden(hidden_size: int, mlp_ratio: float) -> int:
+    r"""
+    SwiGLU intermediate width used by the official blocks.
+
+    Args:
+        hidden_size (`int`):
+            Transformer width.
+        mlp_ratio (`float`):
+            Expansion ratio before the 2/3 SwiGLU reduction.
+
+    Returns:
+        `int`: Intermediate width, rounded up to a multiple of 32.
+    """
+    mlp_hidden = int(hidden_size * float(mlp_ratio))
+    hidden = int((2.0 / 3.0) * mlp_hidden)
+    return ((hidden + 31) // 32) * 32
+
+
+class _StandardMLP(nn.Module):
+    def __init__(self, hidden_size: int, mlp_hidden_dim: int):
+        super().__init__()
+        self.fc1 = _TorchLinear(hidden_size, mlp_hidden_dim, bias=True)
+        self.fc2 = _TorchLinear(mlp_hidden_dim, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        GELU MLP.
+
+        Args:
+            x (`torch.Tensor`):
+                Input `(..., hidden_size)`.
+
+        Returns:
+            `torch.Tensor`: Output `(..., hidden_size)`.
+        """
+        return self.fc2(nn.functional.gelu(self.fc1(x), approximate="none"))
+
+
+class _SwiGLUFFN(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.w1 = _TorchLinear(hidden_size, intermediate_size, bias=True)
+        self.w3 = _TorchLinear(hidden_size, intermediate_size, bias=True)
+        self.w2 = _TorchLinear(intermediate_size, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        SwiGLU feed-forward.
+
+        Args:
+            x (`torch.Tensor`):
+                Input `(..., hidden_size)`.
+
+        Returns:
+            `torch.Tensor`: Output `(..., hidden_size)`.
+        """
+        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
+
+
+class _Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        qk_norm: bool,
+        use_rmsnorm: bool,
+        use_rope: bool,
+        attn_fp32: bool,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.dim = int(dim)
+        self.num_heads = int(num_heads)
+        self.use_rope = bool(use_rope)
+        self.attn_fp32 = bool(attn_fp32)
+        self.qkv = _TorchLinear(self.dim, self.dim * 3, bias=True)
+        head_dim = self.dim // self.num_heads
+        if qk_norm:
+            if use_rmsnorm:
+                self.q_norm = _RMSNorm(head_dim)
+                self.k_norm = _RMSNorm(head_dim)
+            else:
+                self.q_norm = nn.LayerNorm(head_dim, eps=1e-6, elementwise_affine=True)
+                self.k_norm = nn.LayerNorm(head_dim, eps=1e-6, elementwise_affine=True)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+        self.proj = _TorchLinear(self.dim, self.dim, bias=True)
+        if use_rope:
+            half_dim = head_dim // 2
+            freqs = 1.0 / (10000 ** (torch.arange(0, half_dim, dtype=torch.float32) / half_dim))
+            positions = torch.arange(max_seq_len, dtype=torch.float32)
+            emb = torch.cat([torch.outer(positions, freqs)] * 2, dim=-1)
+            self.register_buffer("_rope_cos", torch.cos(emb)[None, :, None, :], persistent=False)
+            self.register_buffer("_rope_sin", torch.sin(emb)[None, :, None, :], persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Multi-head self-attention. Dropout is off, matching inference (`deterministic=True`).
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Attention output `(batch, tokens, hidden)`.
+        """
+        batch, tokens, _channels = x.shape
+        head_dim = self.dim // self.num_heads
+        qkv = self.qkv(x).reshape(batch, tokens, 3, self.num_heads, head_dim)
+        query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        if self.q_norm is not None:
+            query = self.q_norm(query)
+            key = self.k_norm(key)
+        if self.use_rope:
+            rope_dtype = torch.float32 if self.attn_fp32 else query.dtype
+            cos = self._rope_cos[:, :tokens].to(dtype=rope_dtype)
+            sin = self._rope_sin[:, :tokens].to(dtype=rope_dtype)
+            half_dim = head_dim // 2
+
+            def _rotate(tensor: torch.Tensor) -> torch.Tensor:
+                first, second = tensor[..., :half_dim], tensor[..., half_dim:]
+                return torch.cat([-second, first], dim=-1)
+
+            query = (query * cos) + (_rotate(query) * sin)
+            key = (key * cos) + (_rotate(key) * sin)
+        if self.attn_fp32:
+            query = query.float()
+            key = key.float()
+            value = value.float()
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+        if self.attn_fp32:
+            with torch.amp.autocast(device_type=query.device.type, enabled=False):
+                out = nn.functional.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        else:
+            out = nn.functional.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        out = out.permute(0, 2, 1, 3).reshape(batch, tokens, self.dim)
+        return self.proj(out)
+
+
+class _DiTBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float,
+        use_qknorm: bool,
+        use_swiglu: bool,
+        use_rmsnorm: bool,
+        use_rope: bool,
+        attn_fp32: bool,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.hidden_size = int(hidden_size)
+        if use_rmsnorm:
+            self.norm1 = _RMSNorm(self.hidden_size)
+            self.norm2 = _RMSNorm(self.hidden_size)
+        else:
+            self.norm1 = nn.LayerNorm(self.hidden_size, eps=1e-6, elementwise_affine=False)
+            self.norm2 = nn.LayerNorm(self.hidden_size, eps=1e-6, elementwise_affine=False)
+        self.attn = _Attention(
+            dim=self.hidden_size,
+            num_heads=num_heads,
+            qk_norm=use_qknorm,
+            use_rmsnorm=use_rmsnorm,
+            use_rope=use_rope,
+            attn_fp32=attn_fp32,
+            max_seq_len=max_seq_len,
+        )
+        if use_swiglu:
+            self.mlp = _SwiGLUFFN(self.hidden_size, _swiglu_hidden(self.hidden_size, mlp_ratio))
+        else:
+            self.mlp = _StandardMLP(self.hidden_size, int(self.hidden_size * float(mlp_ratio)))
+        self.adaLN_mod = nn.Sequential(
+            nn.SiLU(),
+            _TorchLinear(self.hidden_size, 6 * self.hidden_size, bias=True, weight_init="zeros"),
+        )
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        One Lightning-DiT block conditioned by adaLN-Zero.
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+            condition (`torch.Tensor`):
+                Per-sample condition `(batch, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Updated tokens, same shape as `x`.
+        """
+        chunks = self.adaLN_mod(condition.float()).to(dtype=x.dtype)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(chunks, 6, dim=1)
+        attended = self.attn(_modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + gate_msa.unsqueeze(1) * attended
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(_modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class _FinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int, use_rmsnorm: bool):
+        super().__init__()
+        self.norm_final = _RMSNorm(hidden_size) if use_rmsnorm else nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+        self.adaLN_mod = nn.Sequential(
+            nn.SiLU(),
+            _TorchLinear(hidden_size, 2 * hidden_size, bias=True, weight_init="zeros"),
+        )
+        self.linear = _TorchLinear(hidden_size, patch_size * patch_size * out_channels, bias=True, weight_init="zeros")
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        Project tokens to patch pixels.
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+            condition (`torch.Tensor`):
+                Per-sample condition `(batch, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Patch pixels `(batch, tokens, patch_size * patch_size * out_channels)`.
+        """
+        shift, scale = torch.chunk(self.adaLN_mod(condition.float()).to(dtype=x.dtype), 2, dim=1)
+        return self.linear(_modulate(self.norm_final(x), shift, scale))
+
+
+class _LightningDiT(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        patch_size: int,
+        in_channels: int,
+        hidden_size: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        out_channels: int,
+        use_qknorm: bool,
+        use_swiglu: bool,
+        use_rope: bool,
+        use_rmsnorm: bool,
+        cond_dim: int,
+        n_cls_tokens: int,
+        attn_fp32: bool,
+    ):
+        super().__init__()
+        self.input_size = int(input_size)
+        self.patch_size = int(patch_size)
+        self.in_channels = int(in_channels)
+        self.hidden_size = int(hidden_size)
+        self.out_channels = int(out_channels)
+        self.n_cls_tokens = int(n_cls_tokens)
+        self.patch_embed = _TorchLinear(self.patch_size * self.patch_size * self.in_channels, self.hidden_size, bias=True)
+        grid = self.input_size // self.patch_size
+        max_seq_len = grid * grid + self.n_cls_tokens
+        self.blocks = nn.ModuleList(
+            [
+                _DiTBlock(
+                    hidden_size=self.hidden_size,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    use_qknorm=use_qknorm,
+                    use_swiglu=use_swiglu,
+                    use_rmsnorm=use_rmsnorm,
+                    use_rope=use_rope,
+                    attn_fp32=attn_fp32,
+                    max_seq_len=max_seq_len,
+                )
+                for _ in range(int(depth))
+            ]
+        )
+        self.final_layer = _FinalLayer(self.hidden_size, self.patch_size, self.out_channels, use_rmsnorm)
+        position = _sincos_2d(self.hidden_size, grid)
+        self.pos_embed = nn.Parameter(torch.as_tensor(position, dtype=torch.float32).unsqueeze(0))
+        if self.n_cls_tokens > 0:
+            self.cls_embed = nn.Parameter(torch.randn(1, self.n_cls_tokens, self.hidden_size) * 0.02)
+            self.c_token_proj = _TorchLinear(int(cond_dim), self.hidden_size, bias=True)
+        else:
+            self.register_parameter("cls_embed", None)
+            self.c_token_proj = None
+
+    def forward(self, hidden_states_bhwc: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        Map a latent grid to a clean latent grid.
+
+        Args:
+            hidden_states_bhwc (`torch.Tensor`):
+                Reference noise `(batch, height, width, channels)`. Height and width must equal
+                `input_size`. The official generator uses this layout.
+            condition (`torch.Tensor`):
+                Class, noise-code, and guidance embedding `(batch, cond_dim)`.
+
+        Returns:
+            `torch.Tensor`: Clean latent `(batch, height, width, channels)`.
+        """
+        batch, height, width, channels = hidden_states_bhwc.shape
+        patch = self.patch_size
+        grid = self.input_size // patch
+        if (height, width, channels) != (self.input_size, self.input_size, self.in_channels):
+            raise ValueError(
+                f"Expected latent {(self.input_size, self.input_size, self.in_channels)} in BHWC, "
+                f"got {(height, width, channels)}."
+            )
+        tokens = hidden_states_bhwc.reshape(batch, grid, patch, grid, patch, channels)
+        tokens = tokens.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch, grid * grid, patch * patch * channels)
+        tokens = self.patch_embed(tokens)
+        tokens = (tokens + self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)).to(dtype=tokens.dtype)
+        if self.n_cls_tokens > 0:
+            class_tokens = self.c_token_proj(condition).unsqueeze(1).repeat(1, self.n_cls_tokens, 1)
+            class_tokens = class_tokens + self.cls_embed.to(device=tokens.device, dtype=tokens.dtype)
+            tokens = torch.cat([class_tokens, tokens], dim=1)
+        for block in self.blocks:
+            tokens = block(tokens, condition)
+        tokens = self.final_layer(tokens, condition)
+        if self.n_cls_tokens > 0:
+            tokens = tokens[:, self.n_cls_tokens :, :]
+        tokens = tokens.reshape(batch, grid, grid, patch, patch, self.out_channels)
+        tokens = tokens.permute(0, 1, 3, 2, 4, 5).contiguous()
+        return tokens.reshape(batch, self.input_size, self.input_size, self.out_channels)
+
+
+class _TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.frequency_embedding_size = int(frequency_embedding_size)
+        self.fc1 = _TorchLinear(self.frequency_embedding_size, hidden_size, bias=True, weight_init="normal")
+        self.fc2 = _TorchLinear(hidden_size, hidden_size, bias=True, weight_init="normal")
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        r"""
+        Embed a scalar guidance value with sinusoidal features.
+
+        Args:
+            timesteps (`torch.Tensor`):
+                Shape `(batch,)`. This is the guidance scale, not a diffusion time.
+
+        Returns:
+            `torch.Tensor`: Embedding `(batch, hidden)`.
+        """
+        timesteps = timesteps.to(dtype=torch.float32)
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(0, half, device=timesteps.device, dtype=torch.float32) / half)
+        arguments = timesteps[:, None] * freqs[None]
+        embedding = torch.cat([torch.cos(arguments), torch.sin(arguments)], dim=-1)
+        if self.frequency_embedding_size % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        hidden = nn.functional.silu(self.fc1(embedding))
+        return self.fc2(hidden)
+
+
+class WFlowTransformer2DModel(ModelMixin, ConfigMixin):
+    r"""
+    One-step ImageNet generator used by W-Flow.
+
+    The network maps unit Gaussian SD-VAE noise plus a class label to a clean latent in one pass.
+    Guidance is a continuous embedding (the official `cfg_scale`, equal to `w + 1` in the paper),
+    not a second unconditional forward. Discrete noise codes (`noise_classes` by `noise_coords`)
+    are added to the class embedding.
+
+    Parameter names match `DitGen` in the official repository, so an `ema_model` state dict loads
+    without renaming.
+
+    Args:
+        sample_size (`int`, defaults to 32):
+            Latent spatial size. Released checkpoints use 32, which decodes to 256 pixels.
+        patch_size (`int`, defaults to 2):
+            Patch size on the latent grid.
+        in_channels (`int`, defaults to 4):
+            Latent channels.
+        out_channels (`int`, defaults to 4):
+            Predicted latent channels.
+        cond_dim (`int`, defaults to 768):
+            Width of the class, noise-code, and guidance embeddings.
+        num_classes (`int`, defaults to 1000):
+            ImageNet class count. There is no null-class slot.
+        noise_classes (`int`, defaults to 64):
+            Cardinality of each discrete noise code.
+        noise_coords (`int`, defaults to 32):
+            Number of independent noise-code embeddings that are summed.
+        n_cls_tokens (`int`, defaults to 16):
+            Learned class tokens prepended to the patch sequence.
+        hidden_size (`int`, defaults to 768):
+            Transformer width.
+        depth (`int`, defaults to 12):
+            Number of blocks.
+        num_attention_heads (`int`, defaults to 12):
+            Attention heads.
+        mlp_ratio (`float`, defaults to 4.0):
+            MLP expansion before the SwiGLU 2/3 reduction.
+        use_qknorm (`bool`, defaults to `True`):
+            QK normalization.
+        use_swiglu (`bool`, defaults to `True`):
+            SwiGLU MLP. Released checkpoints turn this on.
+        use_rope (`bool`, defaults to `True`):
+            1-D rotary embeddings over the token sequence, including class tokens.
+        use_rmsnorm (`bool`, defaults to `True`):
+            RMSNorm instead of LayerNorm.
+        use_bf16 (`bool`, defaults to `True`):
+            Autocast the forward to bfloat16 on CUDA, matching the official sampler.
+        attn_fp32 (`bool`, defaults to `False`):
+            Run attention in float32. The ablation checkpoint sets this to `True`.
+        guidance_scale (`float`, defaults to 1.0):
+            Default embedded guidance. The pipeline reads this when the caller omits a scale.
+        default_num_inference_steps (`int`, defaults to 1):
+            Sampling is one network evaluation.
+        model_type (`str`, defaults to `"W-Flow-B"`):
+            Variant name stored in the config.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        cond_dim: int = 768,
+        num_classes: int = 1000,
+        noise_classes: int = 64,
+        noise_coords: int = 32,
+        n_cls_tokens: int = 16,
+        hidden_size: int = 768,
+        depth: int = 12,
+        num_attention_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        use_qknorm: bool = True,
+        use_swiglu: bool = True,
+        use_rope: bool = True,
+        use_rmsnorm: bool = True,
+        use_bf16: bool = True,
+        attn_fp32: bool = False,
+        guidance_scale: float = 1.0,
+        default_num_inference_steps: int = 1,
+        model_type: str = "W-Flow-B",
+    ):
+        super().__init__()
+        self.class_embed = nn.Embedding(int(num_classes), int(cond_dim))
+        nn.init.normal_(self.class_embed.weight, std=0.02)
+        self.noise_embeds = nn.ModuleList()
+        for _ in range(int(noise_coords)):
+            embedding = nn.Embedding(int(noise_classes), int(cond_dim))
+            nn.init.normal_(embedding.weight, std=0.02)
+            self.noise_embeds.append(embedding)
+        self.cfg_embedder = _TimestepEmbedder(int(cond_dim))
+        self.cfg_norm = _RMSNorm(int(cond_dim))
+        self.model = _LightningDiT(
+            input_size=int(sample_size),
+            patch_size=int(patch_size),
+            in_channels=int(in_channels),
+            hidden_size=int(hidden_size),
+            depth=int(depth),
+            num_heads=int(num_attention_heads),
+            mlp_ratio=float(mlp_ratio),
+            out_channels=int(out_channels),
+            use_qknorm=bool(use_qknorm),
+            use_swiglu=bool(use_swiglu),
+            use_rope=bool(use_rope),
+            use_rmsnorm=bool(use_rmsnorm),
+            cond_dim=int(cond_dim),
+            n_cls_tokens=int(n_cls_tokens),
+            attn_fp32=bool(attn_fp32),
+        )
+
+    def condition(
+        self,
+        class_labels: torch.Tensor,
+        guidance_scale: float | torch.Tensor,
+        noise_labels: torch.Tensor,
+    ) -> torch.Tensor:
+        r"""
+        Build the conditioning vector from class, noise codes, and guidance.
+
+        Args:
+            class_labels (`torch.Tensor`):
+                Integer class ids `(batch,)`, in `[0, num_classes)`.
+            guidance_scale (`float` or `torch.Tensor`):
+                Official CFG value `(w + 1)`. A float is broadcast across the batch. A tensor
+                must be `(batch,)` or a scalar.
+            noise_labels (`torch.Tensor`):
+                Integer codes `(batch, noise_coords)` in `[0, noise_classes)`.
+
+        Returns:
+            `torch.Tensor`: Condition `(batch, cond_dim)`.
+        """
+        batch = class_labels.shape[0]
+        cond = self.class_embed(class_labels)
+        for index, embedding in enumerate(self.noise_embeds):
+            cond = cond + embedding(noise_labels[:, index])
+        if isinstance(guidance_scale, (float, int)):
+            scale = torch.full((batch,), float(guidance_scale), device=class_labels.device, dtype=torch.float32)
+        else:
+            scale = torch.as_tensor(guidance_scale, device=class_labels.device, dtype=torch.float32)
+            if scale.ndim == 0:
+                scale = scale[None].repeat(batch)
+        embedded = self.cfg_norm(self.cfg_embedder(scale.float()).float())
+        return cond + embedded.to(dtype=cond.dtype) * 0.02
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        class_labels: torch.Tensor,
+        guidance_scale: float | torch.Tensor = 1.0,
+        noise_labels: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ) -> WFlowTransformer2DModelOutput | tuple[torch.Tensor]:
+        r"""
+        Predict a clean latent from reference noise.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                Gaussian reference noise `(batch, channels, height, width)` in BCHW. Height and
+                width must equal `config.sample_size`.
+            class_labels (`torch.Tensor`):
+                Integer ImageNet ids `(batch,)`.
+            guidance_scale (`float` or `torch.Tensor`, defaults to 1.0):
+                Embedded guidance. This is the official `cfg_scale` (`w + 1` in the paper).
+            noise_labels (`torch.Tensor`, *optional*):
+                Discrete noise codes `(batch, noise_coords)`. Required. The pipeline samples these
+                with the same generator as `hidden_states`.
+            return_dict (`bool`, defaults to `True`):
+                Return [`WFlowTransformer2DModelOutput`] when `True`, otherwise a one-tuple.
+
+        Returns:
+            [`WFlowTransformer2DModelOutput`] or `tuple`:
+                `sample` is the clean latent `(batch, channels, height, width)`.
+        """
+        if noise_labels is None:
+            raise ValueError("`noise_labels` is required and must have shape (batch, noise_coords).")
+        if hidden_states.ndim != 4:
+            raise ValueError(f"`hidden_states` must be BCHW, got {tuple(hidden_states.shape)}.")
+        hidden_bhwc = hidden_states.permute(0, 2, 3, 1).contiguous()
+        use_autocast = bool(self.config.use_bf16) and hidden_states.device.type == "cuda"
+        with torch.amp.autocast(device_type=hidden_states.device.type, dtype=torch.bfloat16, enabled=use_autocast):
+            cond = self.condition(class_labels, guidance_scale, noise_labels)
+            sample_bhwc = self.model(hidden_bhwc, cond)
+        sample = sample_bhwc.permute(0, 3, 1, 2).contiguous()
+        if not return_dict:
+            return (sample,)
+        return WFlowTransformer2DModelOutput(sample=sample)

W-Flow-L/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

W-Flow-L/vae/diffusion_pytorch_model.safetensors

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

W-Flow-XL/model_index.json

@@ -0,0 +1,1021 @@
+{
+  "_class_name": [
+    "pipeline",
+    "WFlowPipeline"
+  ],
+  "_diffusers_version": "0.39.0",
+  "scheduler": [
+    "scheduling_wflow",
+    "WFlowScheduler"
+  ],
+  "transformer": [
+    "transformer_wflow",
+    "WFlowTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ],
+  "id2label": {
+    "0": "tench, Tinca tinca",
+    "1": "goldfish, Carassius auratus",
+    "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
+    "3": "tiger shark, Galeocerdo cuvieri",
+    "4": "hammerhead, hammerhead shark",
+    "5": "electric ray, crampfish, numbfish, torpedo",
+    "6": "stingray",
+    "7": "cock",
+    "8": "hen",
+    "9": "ostrich, Struthio camelus",
+    "10": "brambling, Fringilla montifringilla",
+    "11": "goldfinch, Carduelis carduelis",
+    "12": "house finch, linnet, Carpodacus mexicanus",
+    "13": "junco, snowbird",
+    "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+    "15": "robin, American robin, Turdus migratorius",
+    "16": "bulbul",
+    "17": "jay",
+    "18": "magpie",
+    "19": "chickadee",
+    "20": "water ouzel, dipper",
+    "21": "kite",
+    "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
+    "23": "vulture",
+    "24": "great grey owl, great gray owl, Strix nebulosa",
+    "25": "European fire salamander, Salamandra salamandra",
+    "26": "common newt, Triturus vulgaris",
+    "27": "eft",
+    "28": "spotted salamander, Ambystoma maculatum",
+    "29": "axolotl, mud puppy, Ambystoma mexicanum",
+    "30": "bullfrog, Rana catesbeiana",
+    "31": "tree frog, tree-frog",
+    "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+    "33": "loggerhead, loggerhead turtle, Caretta caretta",
+    "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
+    "35": "mud turtle",
+    "36": "terrapin",
+    "37": "box turtle, box tortoise",
+    "38": "banded gecko",
+    "39": "common iguana, iguana, Iguana iguana",
+    "40": "American chameleon, anole, Anolis carolinensis",
+    "41": "whiptail, whiptail lizard",
+    "42": "agama",
+    "43": "frilled lizard, Chlamydosaurus kingi",
+    "44": "alligator lizard",
+    "45": "Gila monster, Heloderma suspectum",
+    "46": "green lizard, Lacerta viridis",
+    "47": "African chameleon, Chamaeleo chamaeleon",
+    "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
+    "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
+    "50": "American alligator, Alligator mississipiensis",
+    "51": "triceratops",
+    "52": "thunder snake, worm snake, Carphophis amoenus",
+    "53": "ringneck snake, ring-necked snake, ring snake",
+    "54": "hognose snake, puff adder, sand viper",
+    "55": "green snake, grass snake",
+    "56": "king snake, kingsnake",
+    "57": "garter snake, grass snake",
+    "58": "water snake",
+    "59": "vine snake",
+    "60": "night snake, Hypsiglena torquata",
+    "61": "boa constrictor, Constrictor constrictor",
+    "62": "rock python, rock snake, Python sebae",
+    "63": "Indian cobra, Naja naja",
+    "64": "green mamba",
+    "65": "sea snake",
+    "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+    "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+    "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
+    "69": "trilobite",
+    "70": "harvestman, daddy longlegs, Phalangium opilio",
+    "71": "scorpion",
+    "72": "black and gold garden spider, Argiope aurantia",
+    "73": "barn spider, Araneus cavaticus",
+    "74": "garden spider, Aranea diademata",
+    "75": "black widow, Latrodectus mactans",
+    "76": "tarantula",
+    "77": "wolf spider, hunting spider",
+    "78": "tick",
+    "79": "centipede",
+    "80": "black grouse",
+    "81": "ptarmigan",
+    "82": "ruffed grouse, partridge, Bonasa umbellus",
+    "83": "prairie chicken, prairie grouse, prairie fowl",
+    "84": "peacock",
+    "85": "quail",
+    "86": "partridge",
+    "87": "African grey, African gray, Psittacus erithacus",
+    "88": "macaw",
+    "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+    "90": "lorikeet",
+    "91": "coucal",
+    "92": "bee eater",
+    "93": "hornbill",
+    "94": "hummingbird",
+    "95": "jacamar",
+    "96": "toucan",
+    "97": "drake",
+    "98": "red-breasted merganser, Mergus serrator",
+    "99": "goose",
+    "100": "black swan, Cygnus atratus",
+    "101": "tusker",
+    "102": "echidna, spiny anteater, anteater",
+    "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",
+    "104": "wallaby, brush kangaroo",
+    "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",
+    "106": "wombat",
+    "107": "jellyfish",
+    "108": "sea anemone, anemone",
+    "109": "brain coral",
+    "110": "flatworm, platyhelminth",
+    "111": "nematode, nematode worm, roundworm",
+    "112": "conch",
+    "113": "snail",
+    "114": "slug",
+    "115": "sea slug, nudibranch",
+    "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore",
+    "117": "chambered nautilus, pearly nautilus, nautilus",
+    "118": "Dungeness crab, Cancer magister",
+    "119": "rock crab, Cancer irroratus",
+    "120": "fiddler crab",
+    "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica",
+    "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus",
+    "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish",
+    "124": "crayfish, crawfish, crawdad, crawdaddy",
+    "125": "hermit crab",
+    "126": "isopod",
+    "127": "white stork, Ciconia ciconia",
+    "128": "black stork, Ciconia nigra",
+    "129": "spoonbill",
+    "130": "flamingo",
+    "131": "little blue heron, Egretta caerulea",
+    "132": "American egret, great white heron, Egretta albus",
+    "133": "bittern",
+    "134": "crane",
+    "135": "limpkin, Aramus pictus",
+    "136": "European gallinule, Porphyrio porphyrio",
+    "137": "American coot, marsh hen, mud hen, water hen, Fulica americana",
+    "138": "bustard",
+    "139": "ruddy turnstone, Arenaria interpres",
+    "140": "red-backed sandpiper, dunlin, Erolia alpina",
+    "141": "redshank, Tringa totanus",
+    "142": "dowitcher",
+    "143": "oystercatcher, oyster catcher",
+    "144": "pelican",
+    "145": "king penguin, Aptenodytes patagonica",
+    "146": "albatross, mollymawk",
+    "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus",
+    "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca",
+    "149": "dugong, Dugong dugon",
+    "150": "sea lion",
+    "151": "Chihuahua",
+    "152": "Japanese spaniel",
+    "153": "Maltese dog, Maltese terrier, Maltese",
+    "154": "Pekinese, Pekingese, Peke",
+    "155": "Shih-Tzu",
+    "156": "Blenheim spaniel",
+    "157": "papillon",
+    "158": "toy terrier",
+    "159": "Rhodesian ridgeback",
+    "160": "Afghan hound, Afghan",
+    "161": "basset, basset hound",
+    "162": "beagle",
+    "163": "bloodhound, sleuthhound",
+    "164": "bluetick",
+    "165": "black-and-tan coonhound",
+    "166": "Walker hound, Walker foxhound",
+    "167": "English foxhound",
+    "168": "redbone",
+    "169": "borzoi, Russian wolfhound",
+    "170": "Irish wolfhound",
+    "171": "Italian greyhound",
+    "172": "whippet",
+    "173": "Ibizan hound, Ibizan Podenco",
+    "174": "Norwegian elkhound, elkhound",
+    "175": "otterhound, otter hound",
+    "176": "Saluki, gazelle hound",
+    "177": "Scottish deerhound, deerhound",
+    "178": "Weimaraner",
+    "179": "Staffordshire bullterrier, Staffordshire bull terrier",
+    "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",
+    "181": "Bedlington terrier",
+    "182": "Border terrier",
+    "183": "Kerry blue terrier",
+    "184": "Irish terrier",
+    "185": "Norfolk terrier",
+    "186": "Norwich terrier",
+    "187": "Yorkshire terrier",
+    "188": "wire-haired fox terrier",
+    "189": "Lakeland terrier",
+    "190": "Sealyham terrier, Sealyham",
+    "191": "Airedale, Airedale terrier",
+    "192": "cairn, cairn terrier",
+    "193": "Australian terrier",
+    "194": "Dandie Dinmont, Dandie Dinmont terrier",
+    "195": "Boston bull, Boston terrier",
+    "196": "miniature schnauzer",
+    "197": "giant schnauzer",
+    "198": "standard schnauzer",
+    "199": "Scotch terrier, Scottish terrier, Scottie",
+    "200": "Tibetan terrier, chrysanthemum dog",
+    "201": "silky terrier, Sydney silky",
+    "202": "soft-coated wheaten terrier",
+    "203": "West Highland white terrier",
+    "204": "Lhasa, Lhasa apso",
+    "205": "flat-coated retriever",
+    "206": "curly-coated retriever",
+    "207": "golden retriever",
+    "208": "Labrador retriever",
+    "209": "Chesapeake Bay retriever",
+    "210": "German short-haired pointer",
+    "211": "vizsla, Hungarian pointer",
+    "212": "English setter",
+    "213": "Irish setter, red setter",
+    "214": "Gordon setter",
+    "215": "Brittany spaniel",
+    "216": "clumber, clumber spaniel",
+    "217": "English springer, English springer spaniel",
+    "218": "Welsh springer spaniel",
+    "219": "cocker spaniel, English cocker spaniel, cocker",
+    "220": "Sussex spaniel",
+    "221": "Irish water spaniel",
+    "222": "kuvasz",
+    "223": "schipperke",
+    "224": "groenendael",
+    "225": "malinois",
+    "226": "briard",
+    "227": "kelpie",
+    "228": "komondor",
+    "229": "Old English sheepdog, bobtail",
+    "230": "Shetland sheepdog, Shetland sheep dog, Shetland",
+    "231": "collie",
+    "232": "Border collie",
+    "233": "Bouvier des Flandres, Bouviers des Flandres",
+    "234": "Rottweiler",
+    "235": "German shepherd, German shepherd dog, German police dog, alsatian",
+    "236": "Doberman, Doberman pinscher",
+    "237": "miniature pinscher",
+    "238": "Greater Swiss Mountain dog",
+    "239": "Bernese mountain dog",
+    "240": "Appenzeller",
+    "241": "EntleBucher",
+    "242": "boxer",
+    "243": "bull mastiff",
+    "244": "Tibetan mastiff",
+    "245": "French bulldog",
+    "246": "Great Dane",
+    "247": "Saint Bernard, St Bernard",
+    "248": "Eskimo dog, husky",
+    "249": "malamute, malemute, Alaskan malamute",
+    "250": "Siberian husky",
+    "251": "dalmatian, coach dog, carriage dog",
+    "252": "affenpinscher, monkey pinscher, monkey dog",
+    "253": "basenji",
+    "254": "pug, pug-dog",
+    "255": "Leonberg",
+    "256": "Newfoundland, Newfoundland dog",
+    "257": "Great Pyrenees",
+    "258": "Samoyed, Samoyede",
+    "259": "Pomeranian",
+    "260": "chow, chow chow",
+    "261": "keeshond",
+    "262": "Brabancon griffon",
+    "263": "Pembroke, Pembroke Welsh corgi",
+    "264": "Cardigan, Cardigan Welsh corgi",
+    "265": "toy poodle",
+    "266": "miniature poodle",
+    "267": "standard poodle",
+    "268": "Mexican hairless",
+    "269": "timber wolf, grey wolf, gray wolf, Canis lupus",
+    "270": "white wolf, Arctic wolf, Canis lupus tundrarum",
+    "271": "red wolf, maned wolf, Canis rufus, Canis niger",
+    "272": "coyote, prairie wolf, brush wolf, Canis latrans",
+    "273": "dingo, warrigal, warragal, Canis dingo",
+    "274": "dhole, Cuon alpinus",
+    "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
+    "276": "hyena, hyaena",
+    "277": "red fox, Vulpes vulpes",
+    "278": "kit fox, Vulpes macrotis",
+    "279": "Arctic fox, white fox, Alopex lagopus",
+    "280": "grey fox, gray fox, Urocyon cinereoargenteus",
+    "281": "tabby, tabby cat",
+    "282": "tiger cat",
+    "283": "Persian cat",
+    "284": "Siamese cat, Siamese",
+    "285": "Egyptian cat",
+    "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor",
+    "287": "lynx, catamount",
+    "288": "leopard, Panthera pardus",
+    "289": "snow leopard, ounce, Panthera uncia",
+    "290": "jaguar, panther, Panthera onca, Felis onca",
+    "291": "lion, king of beasts, Panthera leo",
+    "292": "tiger, Panthera tigris",
+    "293": "cheetah, chetah, Acinonyx jubatus",
+    "294": "brown bear, bruin, Ursus arctos",
+    "295": "American black bear, black bear, Ursus americanus, Euarctos americanus",
+    "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus",
+    "297": "sloth bear, Melursus ursinus, Ursus ursinus",
+    "298": "mongoose",
+    "299": "meerkat, mierkat",
+    "300": "tiger beetle",
+    "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle",
+    "302": "ground beetle, carabid beetle",
+    "303": "long-horned beetle, longicorn, longicorn beetle",
+    "304": "leaf beetle, chrysomelid",
+    "305": "dung beetle",
+    "306": "rhinoceros beetle",
+    "307": "weevil",
+    "308": "fly",
+    "309": "bee",
+    "310": "ant, emmet, pismire",
+    "311": "grasshopper, hopper",
+    "312": "cricket",
+    "313": "walking stick, walkingstick, stick insect",
+    "314": "cockroach, roach",
+    "315": "mantis, mantid",
+    "316": "cicada, cicala",
+    "317": "leafhopper",
+    "318": "lacewing, lacewing fly",
+    "319": "dragonfly, darning needle, devils darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",
+    "320": "damselfly",
+    "321": "admiral",
+    "322": "ringlet, ringlet butterfly",
+    "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus",
+    "324": "cabbage butterfly",
+    "325": "sulphur butterfly, sulfur butterfly",
+    "326": "lycaenid, lycaenid butterfly",
+    "327": "starfish, sea star",
+    "328": "sea urchin",
+    "329": "sea cucumber, holothurian",
+    "330": "wood rabbit, cottontail, cottontail rabbit",
+    "331": "hare",
+    "332": "Angora, Angora rabbit",
+    "333": "hamster",
+    "334": "porcupine, hedgehog",
+    "335": "fox squirrel, eastern fox squirrel, Sciurus niger",
+    "336": "marmot",
+    "337": "beaver",
+    "338": "guinea pig, Cavia cobaya",
+    "339": "sorrel",
+    "340": "zebra",
+    "341": "hog, pig, grunter, squealer, Sus scrofa",
+    "342": "wild boar, boar, Sus scrofa",
+    "343": "warthog",
+    "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius",
+    "345": "ox",
+    "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis",
+    "347": "bison",
+    "348": "ram, tup",
+    "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis",
+    "350": "ibex, Capra ibex",
+    "351": "hartebeest",
+    "352": "impala, Aepyceros melampus",
+    "353": "gazelle",
+    "354": "Arabian camel, dromedary, Camelus dromedarius",
+    "355": "llama",
+    "356": "weasel",
+    "357": "mink",
+    "358": "polecat, fitch, foulmart, foumart, Mustela putorius",
+    "359": "black-footed ferret, ferret, Mustela nigripes",
+    "360": "otter",
+    "361": "skunk, polecat, wood pussy",
+    "362": "badger",
+    "363": "armadillo",
+    "364": "three-toed sloth, ai, Bradypus tridactylus",
+    "365": "orangutan, orang, orangutang, Pongo pygmaeus",
+    "366": "gorilla, Gorilla gorilla",
+    "367": "chimpanzee, chimp, Pan troglodytes",
+    "368": "gibbon, Hylobates lar",
+    "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus",
+    "370": "guenon, guenon monkey",
+    "371": "patas, hussar monkey, Erythrocebus patas",
+    "372": "baboon",
+    "373": "macaque",
+    "374": "langur",
+    "375": "colobus, colobus monkey",
+    "376": "proboscis monkey, Nasalis larvatus",
+    "377": "marmoset",
+    "378": "capuchin, ringtail, Cebus capucinus",
+    "379": "howler monkey, howler",
+    "380": "titi, titi monkey",
+    "381": "spider monkey, Ateles geoffroyi",
+    "382": "squirrel monkey, Saimiri sciureus",
+    "383": "Madagascar cat, ring-tailed lemur, Lemur catta",
+    "384": "indri, indris, Indri indri, Indri brevicaudatus",
+    "385": "Indian elephant, Elephas maximus",
+    "386": "African elephant, Loxodonta africana",
+    "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens",
+    "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca",
+    "389": "barracouta, snoek",
+    "390": "eel",
+    "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch",
+    "392": "rock beauty, Holocanthus tricolor",
+    "393": "anemone fish",
+    "394": "sturgeon",
+    "395": "gar, garfish, garpike, billfish, Lepisosteus osseus",
+    "396": "lionfish",
+    "397": "puffer, pufferfish, blowfish, globefish",
+    "398": "abacus",
+    "399": "abaya",
+    "400": "academic gown, academic robe, judge robe",
+    "401": "accordion, piano accordion, squeeze box",
+    "402": "acoustic guitar",
+    "403": "aircraft carrier, carrier, flattop, attack aircraft carrier",
+    "404": "airliner",
+    "405": "airship, dirigible",
+    "406": "altar",
+    "407": "ambulance",
+    "408": "amphibian, amphibious vehicle",
+    "409": "analog clock",
+    "410": "apiary, bee house",
+    "411": "apron",
+    "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin",
+    "413": "assault rifle, assault gun",
+    "414": "backpack, back pack, knapsack, packsack, rucksack, haversack",
+    "415": "bakery, bakeshop, bakehouse",
+    "416": "balance beam, beam",
+    "417": "balloon",
+    "418": "ballpoint, ballpoint pen, ballpen, Biro",
+    "419": "Band Aid",
+    "420": "banjo",
+    "421": "bannister, banister, balustrade, balusters, handrail",
+    "422": "barbell",
+    "423": "barber chair",
+    "424": "barbershop",
+    "425": "barn",
+    "426": "barometer",
+    "427": "barrel, cask",
+    "428": "barrow, garden cart, lawn cart, wheelbarrow",
+    "429": "baseball",
+    "430": "basketball",
+    "431": "bassinet",
+    "432": "bassoon",
+    "433": "bathing cap, swimming cap",
+    "434": "bath towel",
+    "435": "bathtub, bathing tub, bath, tub",
+    "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon",
+    "437": "beacon, lighthouse, beacon light, pharos",
+    "438": "beaker",
+    "439": "bearskin, busby, shako",
+    "440": "beer bottle",
+    "441": "beer glass",
+    "442": "bell cote, bell cot",
+    "443": "bib",
+    "444": "bicycle-built-for-two, tandem bicycle, tandem",
+    "445": "bikini, two-piece",
+    "446": "binder, ring-binder",
+    "447": "binoculars, field glasses, opera glasses",
+    "448": "birdhouse",
+    "449": "boathouse",
+    "450": "bobsled, bobsleigh, bob",
+    "451": "bolo tie, bolo, bola tie, bola",
+    "452": "bonnet, poke bonnet",
+    "453": "bookcase",
+    "454": "bookshop, bookstore, bookstall",
+    "455": "bottlecap",
+    "456": "bow",
+    "457": "bow tie, bow-tie, bowtie",
+    "458": "brass, memorial tablet, plaque",
+    "459": "brassiere, bra, bandeau",
+    "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty",
+    "461": "breastplate, aegis, egis",
+    "462": "broom",
+    "463": "bucket, pail",
+    "464": "buckle",
+    "465": "bulletproof vest",
+    "466": "bullet train, bullet",
+    "467": "butcher shop, meat market",
+    "468": "cab, hack, taxi, taxicab",
+    "469": "caldron, cauldron",
+    "470": "candle, taper, wax light",
+    "471": "cannon",
+    "472": "canoe",
+    "473": "can opener, tin opener",
+    "474": "cardigan",
+    "475": "car mirror",
+    "476": "carousel, carrousel, merry-go-round, roundabout, whirligig",
+    "477": "carpenters kit, tool kit",
+    "478": "carton",
+    "479": "car wheel",
+    "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM",
+    "481": "cassette",
+    "482": "cassette player",
+    "483": "castle",
+    "484": "catamaran",
+    "485": "CD player",
+    "486": "cello, violoncello",
+    "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone",
+    "488": "chain",
+    "489": "chainlink fence",
+    "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour",
+    "491": "chain saw, chainsaw",
+    "492": "chest",
+    "493": "chiffonier, commode",
+    "494": "chime, bell, gong",
+    "495": "china cabinet, china closet",
+    "496": "Christmas stocking",
+    "497": "church, church building",
+    "498": "cinema, movie theater, movie theatre, movie house, picture palace",
+    "499": "cleaver, meat cleaver, chopper",
+    "500": "cliff dwelling",
+    "501": "cloak",
+    "502": "clog, geta, patten, sabot",
+    "503": "cocktail shaker",
+    "504": "coffee mug",
+    "505": "coffeepot",
+    "506": "coil, spiral, volute, whorl, helix",
+    "507": "combination lock",
+    "508": "computer keyboard, keypad",
+    "509": "confectionery, confectionary, candy store",
+    "510": "container ship, containership, container vessel",
+    "511": "convertible",
+    "512": "corkscrew, bottle screw",
+    "513": "cornet, horn, trumpet, trump",
+    "514": "cowboy boot",
+    "515": "cowboy hat, ten-gallon hat",
+    "516": "cradle",
+    "517": "crane",
+    "518": "crash helmet",
+    "519": "crate",
+    "520": "crib, cot",
+    "521": "Crock Pot",
+    "522": "croquet ball",
+    "523": "crutch",
+    "524": "cuirass",
+    "525": "dam, dike, dyke",
+    "526": "desk",
+    "527": "desktop computer",
+    "528": "dial telephone, dial phone",
+    "529": "diaper, nappy, napkin",
+    "530": "digital clock",
+    "531": "digital watch",
+    "532": "dining table, board",
+    "533": "dishrag, dishcloth",
+    "534": "dishwasher, dish washer, dishwashing machine",
+    "535": "disk brake, disc brake",
+    "536": "dock, dockage, docking facility",
+    "537": "dogsled, dog sled, dog sleigh",
+    "538": "dome",
+    "539": "doormat, welcome mat",
+    "540": "drilling platform, offshore rig",
+    "541": "drum, membranophone, tympan",
+    "542": "drumstick",
+    "543": "dumbbell",
+    "544": "Dutch oven",
+    "545": "electric fan, blower",
+    "546": "electric guitar",
+    "547": "electric locomotive",
+    "548": "entertainment center",
+    "549": "envelope",
+    "550": "espresso maker",
+    "551": "face powder",
+    "552": "feather boa, boa",
+    "553": "file, file cabinet, filing cabinet",
+    "554": "fireboat",
+    "555": "fire engine, fire truck",
+    "556": "fire screen, fireguard",
+    "557": "flagpole, flagstaff",
+    "558": "flute, transverse flute",
+    "559": "folding chair",
+    "560": "football helmet",
+    "561": "forklift",
+    "562": "fountain",
+    "563": "fountain pen",
+    "564": "four-poster",
+    "565": "freight car",
+    "566": "French horn, horn",
+    "567": "frying pan, frypan, skillet",
+    "568": "fur coat",
+    "569": "garbage truck, dustcart",
+    "570": "gasmask, respirator, gas helmet",
+    "571": "gas pump, gasoline pump, petrol pump, island dispenser",
+    "572": "goblet",
+    "573": "go-kart",
+    "574": "golf ball",
+    "575": "golfcart, golf cart",
+    "576": "gondola",
+    "577": "gong, tam-tam",
+    "578": "gown",
+    "579": "grand piano, grand",
+    "580": "greenhouse, nursery, glasshouse",
+    "581": "grille, radiator grille",
+    "582": "grocery store, grocery, food market, market",
+    "583": "guillotine",
+    "584": "hair slide",
+    "585": "hair spray",
+    "586": "half track",
+    "587": "hammer",
+    "588": "hamper",
+    "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier",
+    "590": "hand-held computer, hand-held microcomputer",
+    "591": "handkerchief, hankie, hanky, hankey",
+    "592": "hard disc, hard disk, fixed disk",
+    "593": "harmonica, mouth organ, harp, mouth harp",
+    "594": "harp",
+    "595": "harvester, reaper",
+    "596": "hatchet",
+    "597": "holster",
+    "598": "home theater, home theatre",
+    "599": "honeycomb",
+    "600": "hook, claw",
+    "601": "hoopskirt, crinoline",
+    "602": "horizontal bar, high bar",
+    "603": "horse cart, horse-cart",
+    "604": "hourglass",
+    "605": "iPod",
+    "606": "iron, smoothing iron",
+    "607": "jack-o-lantern",
+    "608": "jean, blue jean, denim",
+    "609": "jeep, landrover",
+    "610": "jersey, T-shirt, tee shirt",
+    "611": "jigsaw puzzle",
+    "612": "jinrikisha, ricksha, rickshaw",
+    "613": "joystick",
+    "614": "kimono",
+    "615": "knee pad",
+    "616": "knot",
+    "617": "lab coat, laboratory coat",
+    "618": "ladle",
+    "619": "lampshade, lamp shade",
+    "620": "laptop, laptop computer",
+    "621": "lawn mower, mower",
+    "622": "lens cap, lens cover",
+    "623": "letter opener, paper knife, paperknife",
+    "624": "library",
+    "625": "lifeboat",
+    "626": "lighter, light, igniter, ignitor",
+    "627": "limousine, limo",
+    "628": "liner, ocean liner",
+    "629": "lipstick, lip rouge",
+    "630": "Loafer",
+    "631": "lotion",
+    "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system",
+    "633": "loupe, jewelers loupe",
+    "634": "lumbermill, sawmill",
+    "635": "magnetic compass",
+    "636": "mailbag, postbag",
+    "637": "mailbox, letter box",
+    "638": "maillot",
+    "639": "maillot, tank suit",
+    "640": "manhole cover",
+    "641": "maraca",
+    "642": "marimba, xylophone",
+    "643": "mask",
+    "644": "matchstick",
+    "645": "maypole",
+    "646": "maze, labyrinth",
+    "647": "measuring cup",
+    "648": "medicine chest, medicine cabinet",
+    "649": "megalith, megalithic structure",
+    "650": "microphone, mike",
+    "651": "microwave, microwave oven",
+    "652": "military uniform",
+    "653": "milk can",
+    "654": "minibus",
+    "655": "miniskirt, mini",
+    "656": "minivan",
+    "657": "missile",
+    "658": "mitten",
+    "659": "mixing bowl",
+    "660": "mobile home, manufactured home",
+    "661": "Model T",
+    "662": "modem",
+    "663": "monastery",
+    "664": "monitor",
+    "665": "moped",
+    "666": "mortar",
+    "667": "mortarboard",
+    "668": "mosque",
+    "669": "mosquito net",
+    "670": "motor scooter, scooter",
+    "671": "mountain bike, all-terrain bike, off-roader",
+    "672": "mountain tent",
+    "673": "mouse, computer mouse",
+    "674": "mousetrap",
+    "675": "moving van",
+    "676": "muzzle",
+    "677": "nail",
+    "678": "neck brace",
+    "679": "necklace",
+    "680": "nipple",
+    "681": "notebook, notebook computer",
+    "682": "obelisk",
+    "683": "oboe, hautboy, hautbois",
+    "684": "ocarina, sweet potato",
+    "685": "odometer, hodometer, mileometer, milometer",
+    "686": "oil filter",
+    "687": "organ, pipe organ",
+    "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO",
+    "689": "overskirt",
+    "690": "oxcart",
+    "691": "oxygen mask",
+    "692": "packet",
+    "693": "paddle, boat paddle",
+    "694": "paddlewheel, paddle wheel",
+    "695": "padlock",
+    "696": "paintbrush",
+    "697": "pajama, pyjama, pjs, jammies",
+    "698": "palace",
+    "699": "panpipe, pandean pipe, syrinx",
+    "700": "paper towel",
+    "701": "parachute, chute",
+    "702": "parallel bars, bars",
+    "703": "park bench",
+    "704": "parking meter",
+    "705": "passenger car, coach, carriage",
+    "706": "patio, terrace",
+    "707": "pay-phone, pay-station",
+    "708": "pedestal, plinth, footstall",
+    "709": "pencil box, pencil case",
+    "710": "pencil sharpener",
+    "711": "perfume, essence",
+    "712": "Petri dish",
+    "713": "photocopier",
+    "714": "pick, plectrum, plectron",
+    "715": "pickelhaube",
+    "716": "picket fence, paling",
+    "717": "pickup, pickup truck",
+    "718": "pier",
+    "719": "piggy bank, penny bank",
+    "720": "pill bottle",
+    "721": "pillow",
+    "722": "ping-pong ball",
+    "723": "pinwheel",
+    "724": "pirate, pirate ship",
+    "725": "pitcher, ewer",
+    "726": "plane, carpenters plane, woodworking plane",
+    "727": "planetarium",
+    "728": "plastic bag",
+    "729": "plate rack",
+    "730": "plow, plough",
+    "731": "plunger, plumbers helper",
+    "732": "Polaroid camera, Polaroid Land camera",
+    "733": "pole",
+    "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria",
+    "735": "poncho",
+    "736": "pool table, billiard table, snooker table",
+    "737": "pop bottle, soda bottle",
+    "738": "pot, flowerpot",
+    "739": "potters wheel",
+    "740": "power drill",
+    "741": "prayer rug, prayer mat",
+    "742": "printer",
+    "743": "prison, prison house",
+    "744": "projectile, missile",
+    "745": "projector",
+    "746": "puck, hockey puck",
+    "747": "punching bag, punch bag, punching ball, punchball",
+    "748": "purse",
+    "749": "quill, quill pen",
+    "750": "quilt, comforter, comfort, puff",
+    "751": "racer, race car, racing car",
+    "752": "racket, racquet",
+    "753": "radiator",
+    "754": "radio, wireless",
+    "755": "radio telescope, radio reflector",
+    "756": "rain barrel",
+    "757": "recreational vehicle, RV, R.V.",
+    "758": "reel",
+    "759": "reflex camera",
+    "760": "refrigerator, icebox",
+    "761": "remote control, remote",
+    "762": "restaurant, eating house, eating place, eatery",
+    "763": "revolver, six-gun, six-shooter",
+    "764": "rifle",
+    "765": "rocking chair, rocker",
+    "766": "rotisserie",
+    "767": "rubber eraser, rubber, pencil eraser",
+    "768": "rugby ball",
+    "769": "rule, ruler",
+    "770": "running shoe",
+    "771": "safe",
+    "772": "safety pin",
+    "773": "saltshaker, salt shaker",
+    "774": "sandal",
+    "775": "sarong",
+    "776": "sax, saxophone",
+    "777": "scabbard",
+    "778": "scale, weighing machine",
+    "779": "school bus",
+    "780": "schooner",
+    "781": "scoreboard",
+    "782": "screen, CRT screen",
+    "783": "screw",
+    "784": "screwdriver",
+    "785": "seat belt, seatbelt",
+    "786": "sewing machine",
+    "787": "shield, buckler",
+    "788": "shoe shop, shoe-shop, shoe store",
+    "789": "shoji",
+    "790": "shopping basket",
+    "791": "shopping cart",
+    "792": "shovel",
+    "793": "shower cap",
+    "794": "shower curtain",
+    "795": "ski",
+    "796": "ski mask",
+    "797": "sleeping bag",
+    "798": "slide rule, slipstick",
+    "799": "sliding door",
+    "800": "slot, one-armed bandit",
+    "801": "snorkel",
+    "802": "snowmobile",
+    "803": "snowplow, snowplough",
+    "804": "soap dispenser",
+    "805": "soccer ball",
+    "806": "sock",
+    "807": "solar dish, solar collector, solar furnace",
+    "808": "sombrero",
+    "809": "soup bowl",
+    "810": "space bar",
+    "811": "space heater",
+    "812": "space shuttle",
+    "813": "spatula",
+    "814": "speedboat",
+    "815": "spider web, spiders web",
+    "816": "spindle",
+    "817": "sports car, sport car",
+    "818": "spotlight, spot",
+    "819": "stage",
+    "820": "steam locomotive",
+    "821": "steel arch bridge",
+    "822": "steel drum",
+    "823": "stethoscope",
+    "824": "stole",
+    "825": "stone wall",
+    "826": "stopwatch, stop watch",
+    "827": "stove",
+    "828": "strainer",
+    "829": "streetcar, tram, tramcar, trolley, trolley car",
+    "830": "stretcher",
+    "831": "studio couch, day bed",
+    "832": "stupa, tope",
+    "833": "submarine, pigboat, sub, U-boat",
+    "834": "suit, suit of clothes",
+    "835": "sundial",
+    "836": "sunglass",
+    "837": "sunglasses, dark glasses, shades",
+    "838": "sunscreen, sunblock, sun blocker",
+    "839": "suspension bridge",
+    "840": "swab, swob, mop",
+    "841": "sweatshirt",
+    "842": "swimming trunks, bathing trunks",
+    "843": "swing",
+    "844": "switch, electric switch, electrical switch",
+    "845": "syringe",
+    "846": "table lamp",
+    "847": "tank, army tank, armored combat vehicle, armoured combat vehicle",
+    "848": "tape player",
+    "849": "teapot",
+    "850": "teddy, teddy bear",
+    "851": "television, television system",
+    "852": "tennis ball",
+    "853": "thatch, thatched roof",
+    "854": "theater curtain, theatre curtain",
+    "855": "thimble",
+    "856": "thresher, thrasher, threshing machine",
+    "857": "throne",
+    "858": "tile roof",
+    "859": "toaster",
+    "860": "tobacco shop, tobacconist shop, tobacconist",
+    "861": "toilet seat",
+    "862": "torch",
+    "863": "totem pole",
+    "864": "tow truck, tow car, wrecker",
+    "865": "toyshop",
+    "866": "tractor",
+    "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi",
+    "868": "tray",
+    "869": "trench coat",
+    "870": "tricycle, trike, velocipede",
+    "871": "trimaran",
+    "872": "tripod",
+    "873": "triumphal arch",
+    "874": "trolleybus, trolley coach, trackless trolley",
+    "875": "trombone",
+    "876": "tub, vat",
+    "877": "turnstile",
+    "878": "typewriter keyboard",
+    "879": "umbrella",
+    "880": "unicycle, monocycle",
+    "881": "upright, upright piano",
+    "882": "vacuum, vacuum cleaner",
+    "883": "vase",
+    "884": "vault",
+    "885": "velvet",
+    "886": "vending machine",
+    "887": "vestment",
+    "888": "viaduct",
+    "889": "violin, fiddle",
+    "890": "volleyball",
+    "891": "waffle iron",
+    "892": "wall clock",
+    "893": "wallet, billfold, notecase, pocketbook",
+    "894": "wardrobe, closet, press",
+    "895": "warplane, military plane",
+    "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin",
+    "897": "washer, automatic washer, washing machine",
+    "898": "water bottle",
+    "899": "water jug",
+    "900": "water tower",
+    "901": "whiskey jug",
+    "902": "whistle",
+    "903": "wig",
+    "904": "window screen",
+    "905": "window shade",
+    "906": "Windsor tie",
+    "907": "wine bottle",
+    "908": "wing",
+    "909": "wok",
+    "910": "wooden spoon",
+    "911": "wool, woolen, woollen",
+    "912": "worm fence, snake fence, snake-rail fence, Virginia fence",
+    "913": "wreck",
+    "914": "yawl",
+    "915": "yurt",
+    "916": "web site, website, internet site, site",
+    "917": "comic book",
+    "918": "crossword puzzle, crossword",
+    "919": "street sign",
+    "920": "traffic light, traffic signal, stoplight",
+    "921": "book jacket, dust cover, dust jacket, dust wrapper",
+    "922": "menu",
+    "923": "plate",
+    "924": "guacamole",
+    "925": "consomme",
+    "926": "hot pot, hotpot",
+    "927": "trifle",
+    "928": "ice cream, icecream",
+    "929": "ice lolly, lolly, lollipop, popsicle",
+    "930": "French loaf",
+    "931": "bagel, beigel",
+    "932": "pretzel",
+    "933": "cheeseburger",
+    "934": "hotdog, hot dog, red hot",
+    "935": "mashed potato",
+    "936": "head cabbage",
+    "937": "broccoli",
+    "938": "cauliflower",
+    "939": "zucchini, courgette",
+    "940": "spaghetti squash",
+    "941": "acorn squash",
+    "942": "butternut squash",
+    "943": "cucumber, cuke",
+    "944": "artichoke, globe artichoke",
+    "945": "bell pepper",
+    "946": "cardoon",
+    "947": "mushroom",
+    "948": "Granny Smith",
+    "949": "strawberry",
+    "950": "orange",
+    "951": "lemon",
+    "952": "fig",
+    "953": "pineapple, ananas",
+    "954": "banana",
+    "955": "jackfruit, jak, jack",
+    "956": "custard apple",
+    "957": "pomegranate",
+    "958": "hay",
+    "959": "carbonara",
+    "960": "chocolate sauce, chocolate syrup",
+    "961": "dough",
+    "962": "meat loaf, meatloaf",
+    "963": "pizza, pizza pie",
+    "964": "potpie",
+    "965": "burrito",
+    "966": "red wine",
+    "967": "espresso",
+    "968": "cup",
+    "969": "eggnog",
+    "970": "alp",
+    "971": "bubble",
+    "972": "cliff, drop, drop-off",
+    "973": "coral reef",
+    "974": "geyser",
+    "975": "lakeside, lakeshore",
+    "976": "promontory, headland, head, foreland",
+    "977": "sandbar, sand bar",
+    "978": "seashore, coast, seacoast, sea-coast",
+    "979": "valley, vale",
+    "980": "volcano",
+    "981": "ballplayer, baseball player",
+    "982": "groom, bridegroom",
+    "983": "scuba diver",
+    "984": "rapeseed",
+    "985": "daisy",
+    "986": "yellow ladys slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",
+    "987": "corn",
+    "988": "acorn",
+    "989": "hip, rose hip, rosehip",
+    "990": "buckeye, horse chestnut, conker",
+    "991": "coral fungus",
+    "992": "agaric",
+    "993": "gyromitra",
+    "994": "stinkhorn, carrion fungus",
+    "995": "earthstar",
+    "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa",
+    "997": "bolete",
+    "998": "ear, spike, capitulum",
+    "999": "toilet tissue, toilet paper, bathroom tissue"
+  }
+}

W-Flow-XL/pipeline.py

@@ -0,0 +1,754 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.util
+import inspect
+import json
+import sys
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import deprecate, is_torch_xla_available, logging, replace_example_docstring
+from PIL import Image
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from diffusers import DiffusionPipeline
+
+        >>> model_dir = Path("W-Flow-diffusers/W-Flow-XL").resolve()
+        >>> pipe = DiffusionPipeline.from_pretrained(
+        ...     str(model_dir),
+        ...     local_files_only=True,
+        ...     custom_pipeline=str(model_dir / "pipeline.py"),
+        ...     trust_remote_code=True,
+        ...     torch_dtype=torch.bfloat16,
+        ... ).to("cuda")
+
+        >>> image = pipe(
+        ...     class_labels="golden retriever",
+        ...     num_inference_steps=1,
+        ...     guidance_scale=1.09,
+        ...     generator=torch.Generator(device="cuda").manual_seed(1),
+        ... ).images[0]
+        >>> image.save("demo.png")
+
+        >>> # Hugging Face Hub style model id: UserID/RepoID
+        >>> # pipe = DiffusionPipeline.from_pretrained(
+        >>> #     "BiliSakura/W-Flow-diffusers",
+        >>> #     subfolder="W-Flow-XL",
+        >>> #     custom_pipeline="pipeline.py",
+        >>> #     trust_remote_code=True,
+        >>> #     torch_dtype=torch.bfloat16,
+        >>> # )
+        ```
+"""
+
+
+class WFlowPipeline(DiffusionPipeline):
+    r"""
+    Class-conditional one-step pipeline for W-Flow.
+
+    W-Flow maps SD-VAE latent noise to a clean 32×32 latent in one forward pass, then decodes it
+    with `stabilityai/sd-vae-ft-mse` to a 256×256 image. Guidance is embedded in the generator
+    (the official `cfg_scale`, which is `w + 1` in the paper). There is no unconditional class
+    and no second forward.
+
+    The pipeline inherits from [`DiffusionPipeline`]. Check the superclass documentation for the
+    generic methods (download, save, device placement, CPU offload).
+
+    Args:
+        transformer ([`ModelMixin`]):
+            [`WFlowTransformer2DModel`] that predicts the clean latent.
+        vae ([`AutoencoderKL`]):
+            SD-VAE decoder. Released checkpoints use `stabilityai/sd-vae-ft-mse` and scaling 0.18215.
+        scheduler ([`WFlowScheduler`]):
+            One-step scheduler. Loaded from `scheduler/`.
+        id2label (`dict[int, str]`, *optional*):
+            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
+            Loaded from `model_index.json` when omitted.
+
+        The default guidance scale lives on `transformer.config.guidance_scale`.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+    _callback_tensor_inputs = ["latents"]
+
+    def __init__(
+        self,
+        transformer: ModelMixin,
+        vae: AutoencoderKL,
+        scheduler: SchedulerMixin,
+        id2label: dict[int, str] | dict[str, str] | None = None,
+    ):
+        super().__init__()
+        if scheduler is None:
+            scheduler = self._default_scheduler()
+        self._reject_karras_scheduler(scheduler)
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+        normalized_labels = self._normalize_id2label(id2label)
+        self.register_to_config(id2label={str(key): value for key, value in normalized_labels.items()})
+        self._id2label = normalized_labels
+        self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = bool(self._id2label)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    def _sampling_default(self, name: str, fallback: float) -> float:
+        r"""
+        Read a sampling default from the transformer config.
+
+        Args:
+            name (`str`):
+                Key on `transformer.config`, such as `guidance_scale`.
+            fallback (`float`):
+                Value used when the transformer config does not define `name`.
+
+        Returns:
+            `float`: The configured value, or `fallback`.
+        """
+        value = getattr(self.transformer.config, name, None)
+        if value is None:
+            return fallback
+        return value
+
+    @staticmethod
+    def _reject_karras_scheduler(scheduler: SchedulerMixin) -> None:
+        r"""
+        Reject multi-step Karras schedulers.
+
+        W-Flow predicts the clean latent in one evaluation and has no diffusion timestep, so a
+        Karras noise schedule cannot be integrated.
+
+        Args:
+            scheduler ([`SchedulerMixin`]):
+                Scheduler attached to the pipeline.
+        """
+        scheduler_name = scheduler.__class__.__name__
+        if scheduler_name in KarrasDiffusionSchedulers.__members__:
+            raise TypeError(
+                "W-Flow is a one-step generator and does not integrate a Karras noise schedule. "
+                f"Received {scheduler_name}. Use WFlowScheduler."
+            )
+
+    @staticmethod
+    def _default_scheduler() -> SchedulerMixin:
+        r"""
+        Build [`WFlowScheduler`] from the variant directory.
+
+        Returns:
+            [`WFlowScheduler`]: One-step sample scheduler.
+        """
+        module_path = Path(__file__).resolve().parent / "scheduler" / "scheduling_wflow.py"
+        spec = importlib.util.spec_from_file_location("scheduling_wflow", module_path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Could not load the W-Flow scheduler from {module_path}.")
+        module = importlib.util.module_from_spec(spec)
+        sys.modules["scheduling_wflow"] = module
+        spec.loader.exec_module(module)
+        return module.WFlowScheduler()
+
+    def _ensure_labels_loaded(self) -> None:
+        r"""
+        Fill `id2label` from `model_index.json` when the constructor did not receive it.
+        """
+        if self._labels_loaded_from_model_index:
+            return
+        loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None))
+        if loaded:
+            self._id2label = loaded
+            self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = True
+
+    @staticmethod
+    def _normalize_id2label(id2label: dict[int, str] | dict[str, str] | None) -> dict[int, str]:
+        r"""
+        Cast label-map keys to integers.
+
+        Args:
+            id2label (`dict`, *optional*):
+                Class id to English label. Keys may be strings, which is how JSON stores them.
+
+        Returns:
+            `dict[int, str]`: Integer-keyed map. Empty when `id2label` is missing.
+        """
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _read_id2label_from_model_index(variant_path: str | None) -> dict[int, str]:
+        r"""
+        Read `id2label` from a local variant directory.
+
+        Args:
+            variant_path (`str`, *optional*):
+                Directory that contains `model_index.json`.
+
+        Returns:
+            `dict[int, str]`: Loaded map, or an empty dict when the file or key is absent.
+        """
+        if not variant_path:
+            return {}
+        model_index_path = Path(variant_path).resolve() / "model_index.json"
+        if not model_index_path.exists():
+            return {}
+        raw = json.loads(model_index_path.read_text(encoding="utf-8"))
+        id2label = raw.get("id2label")
+        if not isinstance(id2label, dict):
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
+    @staticmethod
+    def _build_label2id(id2label: dict[int, str]) -> dict[str, int]:
+        r"""
+        Invert `id2label`, splitting comma-separated synonyms.
+
+        Args:
+            id2label (`dict[int, str]`):
+                Class id to one or more English names.
+
+        Returns:
+            `dict[str, int]`: Synonym to class id, sorted by synonym.
+        """
+        label2id: dict[str, int] = {}
+        for class_id, value in id2label.items():
+            for synonym in value.split(","):
+                synonym = synonym.strip()
+                if synonym:
+                    label2id[synonym] = int(class_id)
+        return dict(sorted(label2id.items()))
+
+    @property
+    def interrupt(self) -> bool:
+        r"""
+        Whether the current call was asked to stop.
+
+        Returns:
+            `bool`: The `_interrupt` flag. `__call__` sets it to `False` at the start. Assign
+            `_interrupt = True` during a call to skip the remaining forward.
+        """
+        return self._interrupt
+
+    @property
+    def id2label(self) -> dict[int, str]:
+        r"""
+        ImageNet id to English label map.
+
+        Returns:
+            `dict[int, str]`: The map loaded from the constructor or `model_index.json`.
+        """
+        self._ensure_labels_loaded()
+        return self._id2label
+
+    def get_label_ids(self, label: str | list[str]) -> list[int]:
+        r"""
+        Map ImageNet English labels to class ids.
+
+        Args:
+            label (`str` or `list[str]`):
+                One label or a list of labels. Each string must match a synonym in `id2label`
+                after stripping whitespace. Synonyms are the comma-separated pieces of each value.
+
+        Returns:
+            `list[int]`: Class ids in the same order as `label`.
+        """
+        self._ensure_labels_loaded()
+        label2id = self.labels
+        if not label2id:
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
+        if isinstance(label, str):
+            label = [label]
+        missing = [item for item in label if item not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
+        return [label2id[item] for item in label]
+
+    def _normalize_class_labels(self, class_labels: int | str | list[int | str]) -> list[int]:
+        r"""
+        Turn a class-label argument into a list of integer ids.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                One id, one English label, or a list of either. A list must not mix types.
+
+        Returns:
+            `list[int]`: One id per requested image.
+        """
+        if isinstance(class_labels, int):
+            return [class_labels]
+        if isinstance(class_labels, str):
+            return self.get_label_ids(class_labels)
+        if len(class_labels) == 0:
+            raise ValueError("`class_labels` must contain at least one class id or label.")
+        if isinstance(class_labels[0], str):
+            if not all(isinstance(item, str) for item in class_labels):
+                raise TypeError("`class_labels` must be all strings or all integers.")
+            return self.get_label_ids(class_labels)
+        if not all(isinstance(item, int) for item in class_labels):
+            raise TypeError("`class_labels` must be all strings or all integers.")
+        return list(class_labels)
+
+    def prepare_extra_step_kwargs(
+        self,
+        generator: torch.Generator | list[torch.Generator] | None,
+        eta: float,
+    ) -> dict[str, Any]:
+        r"""
+        Extra arguments forwarded to `scheduler.step`, depending on what that method accepts.
+
+        Args:
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Generator passed through when the scheduler step samples noise.
+            eta (`float`):
+                DDIM eta in `[0, 1]`. Ignored by [`WFlowScheduler`], whose `step` has no `eta` argument.
+
+        Returns:
+            `dict`: Keyword arguments for `scheduler.step`.
+        """
+        extra_step_kwargs: dict[str, Any] = {}
+        step_params = set(inspect.signature(self.scheduler.step).parameters.keys())
+        if "eta" in step_params:
+            extra_step_kwargs["eta"] = eta
+        if "generator" in step_params:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        class_labels: int | str | list[int | str],
+        height: int,
+        width: int,
+        num_inference_steps: int,
+        guidance_scale: float,
+        noise_scale: float,
+        output_type: str,
+        callback_steps: int | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+    ) -> None:
+        r"""
+        Reject arguments that the generator cannot run.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                Class ids or English labels. A list must be non-empty.
+            height (`int`):
+                Output height in pixels. Must be the pretrained resolution, 256.
+            width (`int`):
+                Output width in pixels. Must be the pretrained resolution, 256.
+            num_inference_steps (`int`):
+                Must be 1. W-Flow evaluates the network once.
+            guidance_scale (`float`):
+                Embedded guidance (`w + 1` in the paper). Must be finite and positive.
+            noise_scale (`float`):
+                Scale of the reference Gaussian, the official `temp`. Must be positive.
+            output_type (`str`):
+                One of `"pil"`, `"np"`, `"pt"`, `"latent"`.
+            callback_steps (`int`, *optional*):
+                Deprecated callback period. When set, must be a positive integer.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names the step callback may read. Each name must be in `_callback_tensor_inputs`.
+        """
+        if isinstance(class_labels, list) and len(class_labels) == 0:
+            raise ValueError("`class_labels` must contain at least one class id or label.")
+        native = int(self.transformer.config.sample_size) * int(self.vae_scale_factor)
+        if height != native or width != native:
+            raise ValueError(
+                f"`height` and `width` must be the pretrained resolution {native}. "
+                f"The patch embedding is fixed to a {self.transformer.config.sample_size} latent. Got {(height, width)}."
+            )
+        if num_inference_steps != 1:
+            raise ValueError(
+                "W-Flow generates in one forward pass. "
+                f"`num_inference_steps` must be 1, got {num_inference_steps}."
+            )
+        if not torch.isfinite(torch.tensor(float(guidance_scale))) or float(guidance_scale) <= 0:
+            raise ValueError(f"`guidance_scale` must be finite and positive, got {guidance_scale}.")
+        if noise_scale <= 0:
+            raise ValueError(f"`noise_scale` must be positive, got {noise_scale}.")
+        if output_type not in {"pil", "np", "pt", "latent"}:
+            raise ValueError(f"Unsupported `output_type` {output_type!r}. Choose from 'pil', 'np', 'pt', 'latent'.")
+        if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps < 1):
+            raise ValueError(f"`callback_steps` must be a positive integer, got {callback_steps}.")
+        if callback_on_step_end_tensor_inputs is not None:
+            unknown = [name for name in callback_on_step_end_tensor_inputs if name not in self._callback_tensor_inputs]
+            if unknown:
+                raise ValueError(
+                    f"`callback_on_step_end_tensor_inputs` contains {unknown}, which are not in "
+                    f"{self._callback_tensor_inputs}."
+                )
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels: int,
+        height: int,
+        width: int,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None,
+        noise_labels: torch.Tensor | None,
+        noise_scale: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Draw reference noise in the official generator order.
+
+        The official sampler draws `(batch, height, width, channels)` first, then integer noise
+        codes, from the same generator. This method follows that order so a seeded generator
+        matches `DitGen.forward`.
+
+        Args:
+            batch_size (`int`):
+                Number of images.
+            num_channels (`int`):
+                Latent channels. Released checkpoints use 4.
+            height (`int`):
+                Latent height. Released checkpoints use 32.
+            width (`int`):
+                Latent width. Released checkpoints use 32.
+            device (`torch.device`):
+                Device of the returned tensors.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG. A list must have one generator per image. Each generator is consumed in the
+                order Gaussian noise, then noise codes.
+            latents (`torch.Tensor`, *optional*):
+                Reference noise `(batch, channels, height, width)`. When set, it is not rescaled.
+            noise_labels (`torch.Tensor`, *optional*):
+                Integer codes `(batch, noise_coords)`. When omitted, they are drawn here.
+            noise_scale (`float`):
+                Multiplier on the Gaussian, the official `temp`. Ignored when `latents` is passed.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`:
+                Reference noise in BCHW float32, and noise codes `(batch, noise_coords)`.
+        """
+        noise_classes = int(self.transformer.config.noise_classes)
+        noise_coords = int(self.transformer.config.noise_coords)
+        if latents is None:
+            if isinstance(generator, list):
+                if len(generator) != batch_size:
+                    raise ValueError(f"Got {len(generator)} generators for a batch of {batch_size}.")
+                pieces = [
+                    torch.randn((1, height, width, num_channels), generator=one, device=device) for one in generator
+                ]
+                drawn = torch.cat(pieces, dim=0)
+            else:
+                drawn = torch.randn(
+                    (batch_size, height, width, num_channels),
+                    generator=generator,
+                    device=device,
+                )
+            latents = drawn.permute(0, 3, 1, 2).contiguous() * float(noise_scale)
+        elif latents.shape != (batch_size, num_channels, height, width):
+            raise ValueError(
+                f"`latents` must have shape {(batch_size, num_channels, height, width)}, got {tuple(latents.shape)}."
+            )
+        else:
+            latents = latents.to(device=device)
+
+        if noise_labels is None:
+            if isinstance(generator, list) and latents is not None:
+                # When latents were supplied, generators have not been consumed yet only if the
+                # caller passed both. Drawing labels still uses those generators.
+                codes = [
+                    torch.randint(0, noise_classes, (1, noise_coords), generator=one, device=device)
+                    for one in generator
+                ]
+                noise_labels = torch.cat(codes, dim=0)
+            else:
+                noise_labels = torch.randint(
+                    0,
+                    max(1, noise_classes),
+                    (batch_size, max(1, noise_coords)),
+                    generator=None if isinstance(generator, list) else generator,
+                    device=device,
+                )
+        noise_labels = noise_labels.to(device=device, dtype=torch.long)
+        if noise_labels.shape != (batch_size, noise_coords):
+            raise ValueError(
+                f"`noise_labels` must have shape {(batch_size, noise_coords)}, got {tuple(noise_labels.shape)}."
+            )
+        return latents.float(), noise_labels
+
+    def _prepare_scheduler(self, num_inference_steps: int, device: torch.device) -> torch.Tensor:
+        r"""
+        Install the one-step schedule and return its timesteps.
+
+        Args:
+            num_inference_steps (`int`):
+                Must be 1.
+            device (`torch.device`):
+                Device of the timestep tensor.
+
+        Returns:
+            `torch.Tensor`: The single timestep.
+        """
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        return self.scheduler.timesteps
+
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        r"""
+        Decode scaled SD-VAE latents to pixels in `[-1, 1]`.
+
+        Args:
+            latents (`torch.Tensor`):
+                Clean latents `(batch, 4, 32, 32)` in the training scale (`* 0.18215`).
+
+        Returns:
+            `torch.Tensor`: Images `(batch, 3, height, width)` in `[-1, 1]`.
+        """
+        scaling = float(getattr(self.vae.config, "scaling_factor", 0.18215))
+        needs_upcast = bool(getattr(self.vae.config, "force_upcast", False)) and self.vae.dtype != torch.float32
+        original_dtype = self.vae.dtype
+        if needs_upcast:
+            self.vae.to(dtype=torch.float32)
+        decoded = self.vae.decode((latents / scaling).to(dtype=self.vae.dtype), return_dict=False)[0]
+        if needs_upcast:
+            self.vae.to(dtype=original_dtype)
+        return decoded
+
+    def _images_from_pixels(self, pixels: torch.Tensor, output_type: str) -> list[Image.Image] | torch.Tensor | list:
+        r"""
+        Convert decoded pixels to the requested output type.
+
+        Args:
+            pixels (`torch.Tensor`):
+                Images `(batch, 3, height, width)` in `[-1, 1]`.
+            output_type (`str`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"`.
+
+        Returns:
+            Images in the requested type.
+        """
+        images = ((pixels.float() + 1) / 2).clamp(0, 1)
+        if output_type == "pt":
+            return images
+        images = (images * 255).round().to(torch.uint8)
+        arrays = images.permute(0, 2, 3, 1).cpu().numpy()
+        if output_type == "np":
+            return arrays
+        return [Image.fromarray(image) for image in arrays]
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        class_labels: int | str | list[int | str],
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        guidance_scale: float | None = None,
+        noise_scale: float = 1.0,
+        eta: float = 0.0,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        noise_labels: torch.Tensor | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Callable[[int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+        **kwargs,
+    ) -> ImagePipelineOutput | tuple:
+        r"""
+        Generate images from ImageNet class labels.
+
+        Args:
+            class_labels (`int`, `str`, or `list`):
+                One ImageNet class id, one English label, or a list of either. English labels are
+                resolved with [`get_label_ids`]. The list length is the batch size.
+            height (`int`, *optional*):
+                Image height in pixels. Defaults to 256 (`sample_size * vae_scale_factor`). Other
+                sizes are rejected because the patch embedding is fixed.
+            width (`int`, *optional*):
+                Image width in pixels. Defaults to 256.
+            num_inference_steps (`int`, *optional*):
+                Must be 1. Defaults to `transformer.config.default_num_inference_steps`.
+            guidance_scale (`float`, *optional*):
+                Official CFG value, equal to `w + 1` in the paper. Defaults to
+                `transformer.config.guidance_scale` (1.19 / 1.14 / 1.09 / 1.5 by variant). This
+                scale is embedded; the pipeline does not run an unconditional branch.
+            noise_scale (`float`, defaults to 1.0):
+                Multiplier on the reference Gaussian. The official sampler calls this `temp`.
+                Ignored when `latents` is passed.
+            eta (`float`, defaults to 0.0):
+                DDIM eta. Ignored by [`WFlowScheduler`].
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG for the reference noise and the discrete noise codes. One generator is consumed
+                in that order, matching the official sampler.
+            latents (`torch.Tensor`, *optional*):
+                Reference noise `(batch, 4, 32, 32)`. Not multiplied by `noise_scale`.
+            noise_labels (`torch.Tensor`, *optional*):
+                Discrete noise codes `(batch, noise_coords)`. Drawn from `generator` when omitted.
+            output_type (`str`, defaults to `"pil"`):
+                `"pil"`, `"np"`, `"pt"` (float RGB in `[0, 1]`), or `"latent"` (scaled SD-VAE latent).
+            return_dict (`bool`, defaults to `True`):
+                Return [`ImagePipelineOutput`] when `True`, otherwise a one-tuple of images.
+            callback_on_step_end (`Callable` or `PipelineCallback`, *optional*):
+                Called as `callback_on_step_end(pipeline, step, timestep, callback_kwargs)` after the
+                scheduler step. Return a dict to replace tensors listed in
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names passed to the step callback. Defaults to `["latents"]`.
+
+        Examples:
+
+        Returns:
+            [`ImagePipelineOutput`] or `tuple`:
+                When `return_dict` is `True`, [`ImagePipelineOutput`] with the images. Otherwise a tuple
+                whose first element is the images.
+        """
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        if kwargs:
+            raise TypeError(f"Unexpected arguments: {sorted(kwargs)}.")
+        if callback_on_step_end_tensor_inputs is None:
+            callback_on_step_end_tensor_inputs = ["latents"]
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        image_size = int(self.transformer.config.sample_size) * int(self.vae_scale_factor)
+        height = image_size if height is None else int(height)
+        width = image_size if width is None else int(width)
+        if num_inference_steps is None:
+            num_inference_steps = int(self._sampling_default("default_num_inference_steps", 1))
+        if guidance_scale is None:
+            guidance_scale = float(self._sampling_default("guidance_scale", 1.0))
+
+        # 1. Check inputs.
+        self.check_inputs(
+            class_labels,
+            height,
+            width,
+            num_inference_steps,
+            guidance_scale,
+            noise_scale,
+            output_type,
+            callback_steps,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = float(guidance_scale)
+        self._interrupt = False
+
+        # 2. Define call parameters.
+        class_label_ids = self._normalize_class_labels(class_labels)
+        batch_size = len(class_label_ids)
+        device = self._execution_device
+        num_classes = int(self.transformer.config.num_classes)
+        if any(class_id < 0 or class_id >= num_classes for class_id in class_label_ids):
+            raise ValueError(f"Class ids must lie in [0, {num_classes}), got {class_label_ids}.")
+
+        # 3. Encode input condition (class ids; guidance is embedded inside the transformer).
+        class_labels_t = torch.tensor(class_label_ids, device=device, dtype=torch.long)
+
+        # 4. Prepare timesteps.
+        scheduled_timesteps = self._prepare_scheduler(num_inference_steps, device)
+
+        # 5. Prepare latent variables. Noise codes are drawn here so one generator matches DitGen.
+        latent_height = height // int(self.vae_scale_factor)
+        latent_width = width // int(self.vae_scale_factor)
+        latents, noise_labels = self.prepare_latents(
+            batch_size,
+            int(self.transformer.config.in_channels),
+            latent_height,
+            latent_width,
+            device,
+            generator,
+            latents,
+            noise_labels,
+            noise_scale,
+        )
+
+        # 6. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_warmup_steps = len(scheduled_timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(scheduled_timesteps)
+
+        # 7. Denoising loop.
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, timestep in enumerate(scheduled_timesteps):
+                if self.interrupt:
+                    continue
+
+                latent_model_input = self.scheduler.scale_model_input(latents, timestep)
+                model_output = self.transformer(
+                    latent_model_input,
+                    class_labels=class_labels_t,
+                    guidance_scale=self._guidance_scale,
+                    noise_labels=noise_labels,
+                    return_dict=False,
+                )[0]
+                latents = self.scheduler.step(
+                    model_output,
+                    timestep,
+                    latents,
+                    **extra_step_kwargs,
+                    return_dict=False,
+                )[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {key: locals()[key] for key in callback_on_step_end_tensor_inputs}
+                    callback_outputs = callback_on_step_end(self, i, timestep, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+
+                if i == len(scheduled_timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, timestep, latents)
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            images = latents
+        else:
+            images = self._images_from_pixels(self._decode_latents(latents), output_type)
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)

W-Flow-XL/scheduler/scheduler_config.json

@@ -0,0 +1,6 @@
+{
+  "_class_name": "WFlowScheduler",
+  "_diffusers_version": "0.39.0",
+  "num_train_timesteps": 1000,
+  "prediction_type": "sample"
+}

W-Flow-XL/scheduler/scheduling_wflow.py

@@ -0,0 +1,152 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class WFlowSchedulerOutput(BaseOutput):
+    r"""
+    Output of [`WFlowScheduler.step`].
+
+    Args:
+        prev_sample (`torch.Tensor`):
+            The model prediction. W-Flow returns the clean latent in one evaluation.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class WFlowScheduler(SchedulerMixin, ConfigMixin):
+    r"""
+    One-step scheduler for W-Flow.
+
+    The transformer predicts the clean latent directly. `set_timesteps` accepts a single evaluation,
+    and `step` returns that prediction. `scale_model_input` leaves the reference noise unchanged.
+    A Karras noise schedule does not apply: the network has no diffusion timestep.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            Stored for config compatibility. W-Flow does not index it.
+        prediction_type (`str`, defaults to `"sample"`):
+            The model output is the clean sample.
+    """
+
+    order = 1
+    _compatibles: list[str] = []
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        prediction_type: str = "sample",
+    ):
+        if num_train_timesteps < 1:
+            raise ValueError(f"num_train_timesteps must be >= 1, got {num_train_timesteps}.")
+        if prediction_type != "sample":
+            raise ValueError(f"WFlowScheduler prediction_type must be 'sample', got {prediction_type!r}.")
+        self.timesteps: torch.Tensor | None = None
+        self.num_inference_steps: int | None = None
+        self._step_index: int | None = None
+
+    @property
+    def init_noise_sigma(self) -> float:
+        r"""
+        Scale of the standard normal used to start sampling.
+
+        Returns:
+            `float`: `1.0`. The pipeline multiplies this by `noise_scale`.
+        """
+        return 1.0
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor) -> torch.Tensor:
+        r"""
+        Return the sample unchanged.
+
+        Args:
+            sample (`torch.Tensor`):
+                Reference noise `(batch, channels, height, width)`.
+            timestep (`float` or `torch.Tensor`):
+                Unused. W-Flow has no diffusion timestep.
+
+        Returns:
+            `torch.Tensor`: `sample`.
+        """
+        del timestep
+        return sample
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: str | torch.device | None = None,
+    ) -> None:
+        r"""
+        Install the single evaluation step.
+
+        Args:
+            num_inference_steps (`int`):
+                Must be 1. W-Flow generates the latent in one forward pass.
+            device (`str` or `torch.device`, *optional*):
+                Device of `self.timesteps`.
+        """
+        if num_inference_steps != 1:
+            raise ValueError(
+                "W-Flow generates in one forward pass. "
+                f"num_inference_steps must be 1, got {num_inference_steps}."
+            )
+        self.num_inference_steps = 1
+        self.timesteps = torch.tensor([0.0], device=device, dtype=torch.float32)
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: float | torch.Tensor,
+        sample: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        return_dict: bool = True,
+    ) -> WFlowSchedulerOutput | tuple[torch.Tensor]:
+        r"""
+        Replace the reference noise with the predicted clean latent.
+
+        Args:
+            model_output (`torch.Tensor`):
+                Clean-latent prediction `(batch, channels, height, width)`.
+            timestep (`float` or `torch.Tensor`):
+                The single scheduled timestep. Accepted for the Diffusers step signature.
+            sample (`torch.Tensor`):
+                Current reference noise. Same shape as `model_output`.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Accepted so pipelines can pass it through. W-Flow does not draw noise in `step`.
+            return_dict (`bool`, defaults to `True`):
+                Return [`WFlowSchedulerOutput`] when `True`, otherwise a one-tuple.
+
+        Returns:
+            [`WFlowSchedulerOutput`] or `tuple`:
+                `prev_sample` is `model_output`.
+        """
+        del timestep, sample, generator
+        if self.timesteps is None:
+            raise ValueError("Call set_timesteps before step.")
+        if self._step_index is None or self._step_index > 0:
+            raise ValueError("WFlowScheduler has already consumed its single step. Call set_timesteps again.")
+        self._step_index = 1
+        if not return_dict:
+            return (model_output,)
+        return WFlowSchedulerOutput(prev_sample=model_output)

W-Flow-XL/transformer/config.json

@@ -0,0 +1,26 @@
+{
+  "_class_name": "WFlowTransformer2DModel",
+  "_diffusers_version": "0.39.0",
+  "attn_fp32": false,
+  "cond_dim": 1152,
+  "default_num_inference_steps": 1,
+  "depth": 28,
+  "guidance_scale": 1.09,
+  "hidden_size": 1152,
+  "in_channels": 4,
+  "mlp_ratio": 4.0,
+  "model_type": "W-Flow-XL",
+  "n_cls_tokens": 16,
+  "noise_classes": 64,
+  "noise_coords": 32,
+  "num_attention_heads": 16,
+  "num_classes": 1000,
+  "out_channels": 4,
+  "patch_size": 2,
+  "sample_size": 32,
+  "use_bf16": true,
+  "use_qknorm": true,
+  "use_rmsnorm": true,
+  "use_rope": true,
+  "use_swiglu": true
+}

W-Flow-XL/transformer/diffusion_pytorch_model.safetensors

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

W-Flow-XL/transformer/transformer_wflow.py

@@ -0,0 +1,720 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class WFlowTransformer2DModelOutput(BaseOutput):
+    r"""
+    Output of [`WFlowTransformer2DModel`].
+
+    Args:
+        sample (`torch.Tensor`):
+            Predicted SD-VAE latent of shape `(batch, channels, height, width)`. Values are in the
+            scaled latent space used at training time (`encoder_mean * 0.18215`).
+    """
+
+    sample: torch.Tensor
+
+
+def _init_linear(linear: nn.Linear, weight_init: str = "xavier_uniform", bias_init: str = "zeros") -> None:
+    r"""
+    Initialize a linear layer the way the official generator does.
+
+    Args:
+        linear (`torch.nn.Linear`):
+            Layer to initialize in place.
+        weight_init (`str`, defaults to `"xavier_uniform"`):
+            `"xavier_uniform"`, `"zeros"`, or `"normal"` (std 0.02).
+        bias_init (`str`, defaults to `"zeros"`):
+            Bias fill. Only `"zeros"` is used by the released architecture.
+    """
+    if weight_init == "xavier_uniform":
+        nn.init.xavier_uniform_(linear.weight)
+    elif weight_init == "zeros":
+        nn.init.zeros_(linear.weight)
+    elif weight_init == "normal":
+        nn.init.normal_(linear.weight, std=0.02)
+    else:
+        raise ValueError(f"weight_init must be xavier_uniform, zeros, or normal, got {weight_init!r}.")
+    if linear.bias is not None:
+        if bias_init != "zeros":
+            raise ValueError(f"bias_init must be 'zeros', got {bias_init!r}.")
+        nn.init.zeros_(linear.bias)
+
+
+class _TorchLinear(nn.Module):
+    r"""
+    Linear wrapper whose parameter names match the official `TorchLinear` (`*.linear.weight`).
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        weight_init: str = "xavier_uniform",
+        bias_init: str = "zeros",
+    ):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        _init_linear(self.linear, weight_init=weight_init, bias_init=bias_init)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Apply the wrapped linear layer.
+
+        Args:
+            x (`torch.Tensor`):
+                Input of shape `(..., in_features)`.
+
+        Returns:
+            `torch.Tensor`: Output of shape `(..., out_features)`.
+        """
+        return self.linear(x)
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(int(dim)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Root-mean-square normalize the last dimension.
+
+        Args:
+            x (`torch.Tensor`):
+                Input of shape `(..., dim)`.
+
+        Returns:
+            `torch.Tensor`: Normalized tensor, same shape and dtype as `x`.
+        """
+        input_dtype = x.dtype
+        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
+        normalized = x * torch.rsqrt(variance + self.eps)
+        return (normalized * self.weight).to(dtype=input_dtype)
+
+
+def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    r"""
+    Apply an adaLN-Zero shift and scale.
+
+    Args:
+        x (`torch.Tensor`):
+            Token tensor `(batch, tokens, hidden)`.
+        shift (`torch.Tensor`):
+            Per-sample shift `(batch, hidden)`.
+        scale (`torch.Tensor`):
+            Per-sample scale `(batch, hidden)`.
+
+    Returns:
+        `torch.Tensor`: `x * (1 + scale) + shift`, broadcast over tokens.
+    """
+    return x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def _sincos_1d(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    r"""
+    Sine-cosine embedding of 1-D positions.
+
+    Args:
+        embed_dim (`int`):
+            Even embedding width.
+        pos (`numpy.ndarray`):
+            Positions, flattened inside the function.
+
+    Returns:
+        `numpy.ndarray`: Array of shape `(positions, embed_dim)`.
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError(f"embed_dim must be even, got {embed_dim}.")
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega = 1.0 / 10000 ** (omega / (embed_dim / 2.0))
+    out = np.einsum("m,d->md", pos.reshape(-1), omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+def _sincos_2d(embed_dim: int, grid_size: int) -> np.ndarray:
+    r"""
+    Sine-cosine embedding of a square grid, matching the official generator.
+
+    Args:
+        embed_dim (`int`):
+            Full embedding width. Half is used for each axis.
+        grid_size (`int`):
+            Tokens along one side of the latent grid.
+
+    Returns:
+        `numpy.ndarray`: Array of shape `(grid_size * grid_size, embed_dim)`.
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0).reshape(2, 1, grid_size, grid_size)
+    half = embed_dim // 2
+    return np.concatenate([_sincos_1d(half, grid[0]), _sincos_1d(half, grid[1])], axis=1)
+
+
+def _swiglu_hidden(hidden_size: int, mlp_ratio: float) -> int:
+    r"""
+    SwiGLU intermediate width used by the official blocks.
+
+    Args:
+        hidden_size (`int`):
+            Transformer width.
+        mlp_ratio (`float`):
+            Expansion ratio before the 2/3 SwiGLU reduction.
+
+    Returns:
+        `int`: Intermediate width, rounded up to a multiple of 32.
+    """
+    mlp_hidden = int(hidden_size * float(mlp_ratio))
+    hidden = int((2.0 / 3.0) * mlp_hidden)
+    return ((hidden + 31) // 32) * 32
+
+
+class _StandardMLP(nn.Module):
+    def __init__(self, hidden_size: int, mlp_hidden_dim: int):
+        super().__init__()
+        self.fc1 = _TorchLinear(hidden_size, mlp_hidden_dim, bias=True)
+        self.fc2 = _TorchLinear(mlp_hidden_dim, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        GELU MLP.
+
+        Args:
+            x (`torch.Tensor`):
+                Input `(..., hidden_size)`.
+
+        Returns:
+            `torch.Tensor`: Output `(..., hidden_size)`.
+        """
+        return self.fc2(nn.functional.gelu(self.fc1(x), approximate="none"))
+
+
+class _SwiGLUFFN(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.w1 = _TorchLinear(hidden_size, intermediate_size, bias=True)
+        self.w3 = _TorchLinear(hidden_size, intermediate_size, bias=True)
+        self.w2 = _TorchLinear(intermediate_size, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        SwiGLU feed-forward.
+
+        Args:
+            x (`torch.Tensor`):
+                Input `(..., hidden_size)`.
+
+        Returns:
+            `torch.Tensor`: Output `(..., hidden_size)`.
+        """
+        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
+
+
+class _Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        qk_norm: bool,
+        use_rmsnorm: bool,
+        use_rope: bool,
+        attn_fp32: bool,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.dim = int(dim)
+        self.num_heads = int(num_heads)
+        self.use_rope = bool(use_rope)
+        self.attn_fp32 = bool(attn_fp32)
+        self.qkv = _TorchLinear(self.dim, self.dim * 3, bias=True)
+        head_dim = self.dim // self.num_heads
+        if qk_norm:
+            if use_rmsnorm:
+                self.q_norm = _RMSNorm(head_dim)
+                self.k_norm = _RMSNorm(head_dim)
+            else:
+                self.q_norm = nn.LayerNorm(head_dim, eps=1e-6, elementwise_affine=True)
+                self.k_norm = nn.LayerNorm(head_dim, eps=1e-6, elementwise_affine=True)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+        self.proj = _TorchLinear(self.dim, self.dim, bias=True)
+        if use_rope:
+            half_dim = head_dim // 2
+            freqs = 1.0 / (10000 ** (torch.arange(0, half_dim, dtype=torch.float32) / half_dim))
+            positions = torch.arange(max_seq_len, dtype=torch.float32)
+            emb = torch.cat([torch.outer(positions, freqs)] * 2, dim=-1)
+            self.register_buffer("_rope_cos", torch.cos(emb)[None, :, None, :], persistent=False)
+            self.register_buffer("_rope_sin", torch.sin(emb)[None, :, None, :], persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Multi-head self-attention. Dropout is off, matching inference (`deterministic=True`).
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Attention output `(batch, tokens, hidden)`.
+        """
+        batch, tokens, _channels = x.shape
+        head_dim = self.dim // self.num_heads
+        qkv = self.qkv(x).reshape(batch, tokens, 3, self.num_heads, head_dim)
+        query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        if self.q_norm is not None:
+            query = self.q_norm(query)
+            key = self.k_norm(key)
+        if self.use_rope:
+            rope_dtype = torch.float32 if self.attn_fp32 else query.dtype
+            cos = self._rope_cos[:, :tokens].to(dtype=rope_dtype)
+            sin = self._rope_sin[:, :tokens].to(dtype=rope_dtype)
+            half_dim = head_dim // 2
+
+            def _rotate(tensor: torch.Tensor) -> torch.Tensor:
+                first, second = tensor[..., :half_dim], tensor[..., half_dim:]
+                return torch.cat([-second, first], dim=-1)
+
+            query = (query * cos) + (_rotate(query) * sin)
+            key = (key * cos) + (_rotate(key) * sin)
+        if self.attn_fp32:
+            query = query.float()
+            key = key.float()
+            value = value.float()
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+        if self.attn_fp32:
+            with torch.amp.autocast(device_type=query.device.type, enabled=False):
+                out = nn.functional.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        else:
+            out = nn.functional.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+        out = out.permute(0, 2, 1, 3).reshape(batch, tokens, self.dim)
+        return self.proj(out)
+
+
+class _DiTBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float,
+        use_qknorm: bool,
+        use_swiglu: bool,
+        use_rmsnorm: bool,
+        use_rope: bool,
+        attn_fp32: bool,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.hidden_size = int(hidden_size)
+        if use_rmsnorm:
+            self.norm1 = _RMSNorm(self.hidden_size)
+            self.norm2 = _RMSNorm(self.hidden_size)
+        else:
+            self.norm1 = nn.LayerNorm(self.hidden_size, eps=1e-6, elementwise_affine=False)
+            self.norm2 = nn.LayerNorm(self.hidden_size, eps=1e-6, elementwise_affine=False)
+        self.attn = _Attention(
+            dim=self.hidden_size,
+            num_heads=num_heads,
+            qk_norm=use_qknorm,
+            use_rmsnorm=use_rmsnorm,
+            use_rope=use_rope,
+            attn_fp32=attn_fp32,
+            max_seq_len=max_seq_len,
+        )
+        if use_swiglu:
+            self.mlp = _SwiGLUFFN(self.hidden_size, _swiglu_hidden(self.hidden_size, mlp_ratio))
+        else:
+            self.mlp = _StandardMLP(self.hidden_size, int(self.hidden_size * float(mlp_ratio)))
+        self.adaLN_mod = nn.Sequential(
+            nn.SiLU(),
+            _TorchLinear(self.hidden_size, 6 * self.hidden_size, bias=True, weight_init="zeros"),
+        )
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        One Lightning-DiT block conditioned by adaLN-Zero.
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+            condition (`torch.Tensor`):
+                Per-sample condition `(batch, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Updated tokens, same shape as `x`.
+        """
+        chunks = self.adaLN_mod(condition.float()).to(dtype=x.dtype)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(chunks, 6, dim=1)
+        attended = self.attn(_modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + gate_msa.unsqueeze(1) * attended
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(_modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class _FinalLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int, use_rmsnorm: bool):
+        super().__init__()
+        self.norm_final = _RMSNorm(hidden_size) if use_rmsnorm else nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+        self.adaLN_mod = nn.Sequential(
+            nn.SiLU(),
+            _TorchLinear(hidden_size, 2 * hidden_size, bias=True, weight_init="zeros"),
+        )
+        self.linear = _TorchLinear(hidden_size, patch_size * patch_size * out_channels, bias=True, weight_init="zeros")
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        Project tokens to patch pixels.
+
+        Args:
+            x (`torch.Tensor`):
+                Tokens `(batch, tokens, hidden)`.
+            condition (`torch.Tensor`):
+                Per-sample condition `(batch, hidden)`.
+
+        Returns:
+            `torch.Tensor`: Patch pixels `(batch, tokens, patch_size * patch_size * out_channels)`.
+        """
+        shift, scale = torch.chunk(self.adaLN_mod(condition.float()).to(dtype=x.dtype), 2, dim=1)
+        return self.linear(_modulate(self.norm_final(x), shift, scale))
+
+
+class _LightningDiT(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        patch_size: int,
+        in_channels: int,
+        hidden_size: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        out_channels: int,
+        use_qknorm: bool,
+        use_swiglu: bool,
+        use_rope: bool,
+        use_rmsnorm: bool,
+        cond_dim: int,
+        n_cls_tokens: int,
+        attn_fp32: bool,
+    ):
+        super().__init__()
+        self.input_size = int(input_size)
+        self.patch_size = int(patch_size)
+        self.in_channels = int(in_channels)
+        self.hidden_size = int(hidden_size)
+        self.out_channels = int(out_channels)
+        self.n_cls_tokens = int(n_cls_tokens)
+        self.patch_embed = _TorchLinear(self.patch_size * self.patch_size * self.in_channels, self.hidden_size, bias=True)
+        grid = self.input_size // self.patch_size
+        max_seq_len = grid * grid + self.n_cls_tokens
+        self.blocks = nn.ModuleList(
+            [
+                _DiTBlock(
+                    hidden_size=self.hidden_size,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    use_qknorm=use_qknorm,
+                    use_swiglu=use_swiglu,
+                    use_rmsnorm=use_rmsnorm,
+                    use_rope=use_rope,
+                    attn_fp32=attn_fp32,
+                    max_seq_len=max_seq_len,
+                )
+                for _ in range(int(depth))
+            ]
+        )
+        self.final_layer = _FinalLayer(self.hidden_size, self.patch_size, self.out_channels, use_rmsnorm)
+        position = _sincos_2d(self.hidden_size, grid)
+        self.pos_embed = nn.Parameter(torch.as_tensor(position, dtype=torch.float32).unsqueeze(0))
+        if self.n_cls_tokens > 0:
+            self.cls_embed = nn.Parameter(torch.randn(1, self.n_cls_tokens, self.hidden_size) * 0.02)
+            self.c_token_proj = _TorchLinear(int(cond_dim), self.hidden_size, bias=True)
+        else:
+            self.register_parameter("cls_embed", None)
+            self.c_token_proj = None
+
+    def forward(self, hidden_states_bhwc: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        r"""
+        Map a latent grid to a clean latent grid.
+
+        Args:
+            hidden_states_bhwc (`torch.Tensor`):
+                Reference noise `(batch, height, width, channels)`. Height and width must equal
+                `input_size`. The official generator uses this layout.
+            condition (`torch.Tensor`):
+                Class, noise-code, and guidance embedding `(batch, cond_dim)`.
+
+        Returns:
+            `torch.Tensor`: Clean latent `(batch, height, width, channels)`.
+        """
+        batch, height, width, channels = hidden_states_bhwc.shape
+        patch = self.patch_size
+        grid = self.input_size // patch
+        if (height, width, channels) != (self.input_size, self.input_size, self.in_channels):
+            raise ValueError(
+                f"Expected latent {(self.input_size, self.input_size, self.in_channels)} in BHWC, "
+                f"got {(height, width, channels)}."
+            )
+        tokens = hidden_states_bhwc.reshape(batch, grid, patch, grid, patch, channels)
+        tokens = tokens.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch, grid * grid, patch * patch * channels)
+        tokens = self.patch_embed(tokens)
+        tokens = (tokens + self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)).to(dtype=tokens.dtype)
+        if self.n_cls_tokens > 0:
+            class_tokens = self.c_token_proj(condition).unsqueeze(1).repeat(1, self.n_cls_tokens, 1)
+            class_tokens = class_tokens + self.cls_embed.to(device=tokens.device, dtype=tokens.dtype)
+            tokens = torch.cat([class_tokens, tokens], dim=1)
+        for block in self.blocks:
+            tokens = block(tokens, condition)
+        tokens = self.final_layer(tokens, condition)
+        if self.n_cls_tokens > 0:
+            tokens = tokens[:, self.n_cls_tokens :, :]
+        tokens = tokens.reshape(batch, grid, grid, patch, patch, self.out_channels)
+        tokens = tokens.permute(0, 1, 3, 2, 4, 5).contiguous()
+        return tokens.reshape(batch, self.input_size, self.input_size, self.out_channels)
+
+
+class _TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.frequency_embedding_size = int(frequency_embedding_size)
+        self.fc1 = _TorchLinear(self.frequency_embedding_size, hidden_size, bias=True, weight_init="normal")
+        self.fc2 = _TorchLinear(hidden_size, hidden_size, bias=True, weight_init="normal")
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        r"""
+        Embed a scalar guidance value with sinusoidal features.
+
+        Args:
+            timesteps (`torch.Tensor`):
+                Shape `(batch,)`. This is the guidance scale, not a diffusion time.
+
+        Returns:
+            `torch.Tensor`: Embedding `(batch, hidden)`.
+        """
+        timesteps = timesteps.to(dtype=torch.float32)
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(0, half, device=timesteps.device, dtype=torch.float32) / half)
+        arguments = timesteps[:, None] * freqs[None]
+        embedding = torch.cat([torch.cos(arguments), torch.sin(arguments)], dim=-1)
+        if self.frequency_embedding_size % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        hidden = nn.functional.silu(self.fc1(embedding))
+        return self.fc2(hidden)
+
+
+class WFlowTransformer2DModel(ModelMixin, ConfigMixin):
+    r"""
+    One-step ImageNet generator used by W-Flow.
+
+    The network maps unit Gaussian SD-VAE noise plus a class label to a clean latent in one pass.
+    Guidance is a continuous embedding (the official `cfg_scale`, equal to `w + 1` in the paper),
+    not a second unconditional forward. Discrete noise codes (`noise_classes` by `noise_coords`)
+    are added to the class embedding.
+
+    Parameter names match `DitGen` in the official repository, so an `ema_model` state dict loads
+    without renaming.
+
+    Args:
+        sample_size (`int`, defaults to 32):
+            Latent spatial size. Released checkpoints use 32, which decodes to 256 pixels.
+        patch_size (`int`, defaults to 2):
+            Patch size on the latent grid.
+        in_channels (`int`, defaults to 4):
+            Latent channels.
+        out_channels (`int`, defaults to 4):
+            Predicted latent channels.
+        cond_dim (`int`, defaults to 768):
+            Width of the class, noise-code, and guidance embeddings.
+        num_classes (`int`, defaults to 1000):
+            ImageNet class count. There is no null-class slot.
+        noise_classes (`int`, defaults to 64):
+            Cardinality of each discrete noise code.
+        noise_coords (`int`, defaults to 32):
+            Number of independent noise-code embeddings that are summed.
+        n_cls_tokens (`int`, defaults to 16):
+            Learned class tokens prepended to the patch sequence.
+        hidden_size (`int`, defaults to 768):
+            Transformer width.
+        depth (`int`, defaults to 12):
+            Number of blocks.
+        num_attention_heads (`int`, defaults to 12):
+            Attention heads.
+        mlp_ratio (`float`, defaults to 4.0):
+            MLP expansion before the SwiGLU 2/3 reduction.
+        use_qknorm (`bool`, defaults to `True`):
+            QK normalization.
+        use_swiglu (`bool`, defaults to `True`):
+            SwiGLU MLP. Released checkpoints turn this on.
+        use_rope (`bool`, defaults to `True`):
+            1-D rotary embeddings over the token sequence, including class tokens.
+        use_rmsnorm (`bool`, defaults to `True`):
+            RMSNorm instead of LayerNorm.
+        use_bf16 (`bool`, defaults to `True`):
+            Autocast the forward to bfloat16 on CUDA, matching the official sampler.
+        attn_fp32 (`bool`, defaults to `False`):
+            Run attention in float32. The ablation checkpoint sets this to `True`.
+        guidance_scale (`float`, defaults to 1.0):
+            Default embedded guidance. The pipeline reads this when the caller omits a scale.
+        default_num_inference_steps (`int`, defaults to 1):
+            Sampling is one network evaluation.
+        model_type (`str`, defaults to `"W-Flow-B"`):
+            Variant name stored in the config.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        cond_dim: int = 768,
+        num_classes: int = 1000,
+        noise_classes: int = 64,
+        noise_coords: int = 32,
+        n_cls_tokens: int = 16,
+        hidden_size: int = 768,
+        depth: int = 12,
+        num_attention_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        use_qknorm: bool = True,
+        use_swiglu: bool = True,
+        use_rope: bool = True,
+        use_rmsnorm: bool = True,
+        use_bf16: bool = True,
+        attn_fp32: bool = False,
+        guidance_scale: float = 1.0,
+        default_num_inference_steps: int = 1,
+        model_type: str = "W-Flow-B",
+    ):
+        super().__init__()
+        self.class_embed = nn.Embedding(int(num_classes), int(cond_dim))
+        nn.init.normal_(self.class_embed.weight, std=0.02)
+        self.noise_embeds = nn.ModuleList()
+        for _ in range(int(noise_coords)):
+            embedding = nn.Embedding(int(noise_classes), int(cond_dim))
+            nn.init.normal_(embedding.weight, std=0.02)
+            self.noise_embeds.append(embedding)
+        self.cfg_embedder = _TimestepEmbedder(int(cond_dim))
+        self.cfg_norm = _RMSNorm(int(cond_dim))
+        self.model = _LightningDiT(
+            input_size=int(sample_size),
+            patch_size=int(patch_size),
+            in_channels=int(in_channels),
+            hidden_size=int(hidden_size),
+            depth=int(depth),
+            num_heads=int(num_attention_heads),
+            mlp_ratio=float(mlp_ratio),
+            out_channels=int(out_channels),
+            use_qknorm=bool(use_qknorm),
+            use_swiglu=bool(use_swiglu),
+            use_rope=bool(use_rope),
+            use_rmsnorm=bool(use_rmsnorm),
+            cond_dim=int(cond_dim),
+            n_cls_tokens=int(n_cls_tokens),
+            attn_fp32=bool(attn_fp32),
+        )
+
+    def condition(
+        self,
+        class_labels: torch.Tensor,
+        guidance_scale: float | torch.Tensor,
+        noise_labels: torch.Tensor,
+    ) -> torch.Tensor:
+        r"""
+        Build the conditioning vector from class, noise codes, and guidance.
+
+        Args:
+            class_labels (`torch.Tensor`):
+                Integer class ids `(batch,)`, in `[0, num_classes)`.
+            guidance_scale (`float` or `torch.Tensor`):
+                Official CFG value `(w + 1)`. A float is broadcast across the batch. A tensor
+                must be `(batch,)` or a scalar.
+            noise_labels (`torch.Tensor`):
+                Integer codes `(batch, noise_coords)` in `[0, noise_classes)`.
+
+        Returns:
+            `torch.Tensor`: Condition `(batch, cond_dim)`.
+        """
+        batch = class_labels.shape[0]
+        cond = self.class_embed(class_labels)
+        for index, embedding in enumerate(self.noise_embeds):
+            cond = cond + embedding(noise_labels[:, index])
+        if isinstance(guidance_scale, (float, int)):
+            scale = torch.full((batch,), float(guidance_scale), device=class_labels.device, dtype=torch.float32)
+        else:
+            scale = torch.as_tensor(guidance_scale, device=class_labels.device, dtype=torch.float32)
+            if scale.ndim == 0:
+                scale = scale[None].repeat(batch)
+        embedded = self.cfg_norm(self.cfg_embedder(scale.float()).float())
+        return cond + embedded.to(dtype=cond.dtype) * 0.02
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        class_labels: torch.Tensor,
+        guidance_scale: float | torch.Tensor = 1.0,
+        noise_labels: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ) -> WFlowTransformer2DModelOutput | tuple[torch.Tensor]:
+        r"""
+        Predict a clean latent from reference noise.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                Gaussian reference noise `(batch, channels, height, width)` in BCHW. Height and
+                width must equal `config.sample_size`.
+            class_labels (`torch.Tensor`):
+                Integer ImageNet ids `(batch,)`.
+            guidance_scale (`float` or `torch.Tensor`, defaults to 1.0):
+                Embedded guidance. This is the official `cfg_scale` (`w + 1` in the paper).
+            noise_labels (`torch.Tensor`, *optional*):
+                Discrete noise codes `(batch, noise_coords)`. Required. The pipeline samples these
+                with the same generator as `hidden_states`.
+            return_dict (`bool`, defaults to `True`):
+                Return [`WFlowTransformer2DModelOutput`] when `True`, otherwise a one-tuple.
+
+        Returns:
+            [`WFlowTransformer2DModelOutput`] or `tuple`:
+                `sample` is the clean latent `(batch, channels, height, width)`.
+        """
+        if noise_labels is None:
+            raise ValueError("`noise_labels` is required and must have shape (batch, noise_coords).")
+        if hidden_states.ndim != 4:
+            raise ValueError(f"`hidden_states` must be BCHW, got {tuple(hidden_states.shape)}.")
+        hidden_bhwc = hidden_states.permute(0, 2, 3, 1).contiguous()
+        use_autocast = bool(self.config.use_bf16) and hidden_states.device.type == "cuda"
+        with torch.amp.autocast(device_type=hidden_states.device.type, dtype=torch.bfloat16, enabled=use_autocast):
+            cond = self.condition(class_labels, guidance_scale, noise_labels)
+            sample_bhwc = self.model(hidden_bhwc, cond)
+        sample = sample_bhwc.permute(0, 3, 1, 2).contiguous()
+        if not return_dict:
+            return (sample,)
+        return WFlowTransformer2DModelOutput(sample=sample)

W-Flow-XL/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

W-Flow-XL/vae/diffusion_pytorch_model.safetensors

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