BiliSakura commited on
Commit
c205191
·
verified ·
1 Parent(s): 9aa9fef

Upload folder using huggingface_hub

Browse files
DiT-MoE-B-8E2A/README.md ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ library_name: diffusers
4
+ pipeline_tag: unconditional-image-generation
5
+ tags:
6
+ - diffusers
7
+ - dit-moe
8
+ - image-generation
9
+ - class-conditional
10
+ - imagenet
11
+ inference: true
12
+ ---
13
+
14
+ # DiT-MoE-B-8E2A
15
+
16
+ Self-contained Diffusers checkpoint for **DiT-B/2** with MoE routing, converted from [`feizhengcong/DiT-MoE`](https://huggingface.co/feizhengcong/DiT-MoE).
17
+
18
+ Each subfolder is a self-contained Diffusers model repo with:
19
+
20
+ - `model_index.json` (includes ImageNet `id2label`)
21
+ - `pipeline.py` (custom `DiTMoEPipeline`)
22
+ - `transformer/transformer_dit_moe.py` and weights
23
+ - `vae/diffusion_pytorch_model.safetensors`
24
+ - `scheduler/scheduler_config.json`
25
+
26
+ ## ImageNet class labels
27
+
28
+ Each variant keeps an English `id2label` map in `model_index.json` (DiT-style).
29
+
30
+ - `pipe.id2label[207]` — `"golden retriever"`
31
+ - `pipe.get_label_ids("golden retriever")` — `[207]`
32
+ - `pipe(class_labels="golden retriever", ...)` — string labels resolved automatically
33
+
34
+ ## Recommended inference (256×256)
35
+
36
+ | Setting | Value |
37
+ | --- | --- |
38
+ | Resolution | 256×256 |
39
+ | Sampler | DDIM |
40
+ | Steps | 50 |
41
+ | CFG scale | 4.0 |
42
+ | Dtype | `bfloat16` (recommended on Ampere+) |
43
+ | VAE | `stabilityai/sd-vae-ft-mse` (bundled under `vae/`) |
44
+
45
+ ```python
46
+ from pathlib import Path
47
+ import torch
48
+ from diffusers import DiffusionPipeline
49
+
50
+ model_dir = Path("./DiT-MoE-B-8E2A").resolve()
51
+ pipe = DiffusionPipeline.from_pretrained(
52
+ str(model_dir),
53
+ local_files_only=True,
54
+ custom_pipeline=str(model_dir / "pipeline.py"),
55
+ trust_remote_code=True,
56
+ torch_dtype=torch.bfloat16,
57
+ )
58
+ pipe.to("cuda")
59
+
60
+ print(pipe.id2label[207])
61
+ print(pipe.get_label_ids("golden retriever"))
62
+
63
+ generator = torch.Generator(device="cuda").manual_seed(42)
64
+ image = pipe(
65
+ class_labels="golden retriever",
66
+ height=256,
67
+ width=256,
68
+ num_inference_steps=50,
69
+ guidance_scale=4.0,
70
+ generator=generator,
71
+ ).images[0]
72
+ image.save("demo.png")
73
+ ```
DiT-MoE-B-8E2A/model_index.json ADDED
@@ -0,0 +1,1021 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": [
3
+ "pipeline",
4
+ "DiTMoEPipeline"
5
+ ],
6
+ "_diffusers_version": "0.36.0",
7
+ "scheduler": [
8
+ "diffusers",
9
+ "DDIMScheduler"
10
+ ],
11
+ "transformer": [
12
+ "transformer_dit_moe",
13
+ "DiTMoETransformer2DModel"
14
+ ],
15
+ "vae": [
16
+ "diffusers",
17
+ "AutoencoderKL"
18
+ ],
19
+ "id2label": {
20
+ "0": "tench, Tinca tinca",
21
+ "1": "goldfish, Carassius auratus",
22
+ "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
23
+ "3": "tiger shark, Galeocerdo cuvieri",
24
+ "4": "hammerhead, hammerhead shark",
25
+ "5": "electric ray, crampfish, numbfish, torpedo",
26
+ "6": "stingray",
27
+ "7": "cock",
28
+ "8": "hen",
29
+ "9": "ostrich, Struthio camelus",
30
+ "10": "brambling, Fringilla montifringilla",
31
+ "11": "goldfinch, Carduelis carduelis",
32
+ "12": "house finch, linnet, Carpodacus mexicanus",
33
+ "13": "junco, snowbird",
34
+ "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
35
+ "15": "robin, American robin, Turdus migratorius",
36
+ "16": "bulbul",
37
+ "17": "jay",
38
+ "18": "magpie",
39
+ "19": "chickadee",
40
+ "20": "water ouzel, dipper",
41
+ "21": "kite",
42
+ "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
43
+ "23": "vulture",
44
+ "24": "great grey owl, great gray owl, Strix nebulosa",
45
+ "25": "European fire salamander, Salamandra salamandra",
46
+ "26": "common newt, Triturus vulgaris",
47
+ "27": "eft",
48
+ "28": "spotted salamander, Ambystoma maculatum",
49
+ "29": "axolotl, mud puppy, Ambystoma mexicanum",
50
+ "30": "bullfrog, Rana catesbeiana",
51
+ "31": "tree frog, tree-frog",
52
+ "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
53
+ "33": "loggerhead, loggerhead turtle, Caretta caretta",
54
+ "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
55
+ "35": "mud turtle",
56
+ "36": "terrapin",
57
+ "37": "box turtle, box tortoise",
58
+ "38": "banded gecko",
59
+ "39": "common iguana, iguana, Iguana iguana",
60
+ "40": "American chameleon, anole, Anolis carolinensis",
61
+ "41": "whiptail, whiptail lizard",
62
+ "42": "agama",
63
+ "43": "frilled lizard, Chlamydosaurus kingi",
64
+ "44": "alligator lizard",
65
+ "45": "Gila monster, Heloderma suspectum",
66
+ "46": "green lizard, Lacerta viridis",
67
+ "47": "African chameleon, Chamaeleo chamaeleon",
68
+ "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
69
+ "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
70
+ "50": "American alligator, Alligator mississipiensis",
71
+ "51": "triceratops",
72
+ "52": "thunder snake, worm snake, Carphophis amoenus",
73
+ "53": "ringneck snake, ring-necked snake, ring snake",
74
+ "54": "hognose snake, puff adder, sand viper",
75
+ "55": "green snake, grass snake",
76
+ "56": "king snake, kingsnake",
77
+ "57": "garter snake, grass snake",
78
+ "58": "water snake",
79
+ "59": "vine snake",
80
+ "60": "night snake, Hypsiglena torquata",
81
+ "61": "boa constrictor, Constrictor constrictor",
82
+ "62": "rock python, rock snake, Python sebae",
83
+ "63": "Indian cobra, Naja naja",
84
+ "64": "green mamba",
85
+ "65": "sea snake",
86
+ "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
87
+ "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
88
+ "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
89
+ "69": "trilobite",
90
+ "70": "harvestman, daddy longlegs, Phalangium opilio",
91
+ "71": "scorpion",
92
+ "72": "black and gold garden spider, Argiope aurantia",
93
+ "73": "barn spider, Araneus cavaticus",
94
+ "74": "garden spider, Aranea diademata",
95
+ "75": "black widow, Latrodectus mactans",
96
+ "76": "tarantula",
97
+ "77": "wolf spider, hunting spider",
98
+ "78": "tick",
99
+ "79": "centipede",
100
+ "80": "black grouse",
101
+ "81": "ptarmigan",
102
+ "82": "ruffed grouse, partridge, Bonasa umbellus",
103
+ "83": "prairie chicken, prairie grouse, prairie fowl",
104
+ "84": "peacock",
105
+ "85": "quail",
106
+ "86": "partridge",
107
+ "87": "African grey, African gray, Psittacus erithacus",
108
+ "88": "macaw",
109
+ "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
110
+ "90": "lorikeet",
111
+ "91": "coucal",
112
+ "92": "bee eater",
113
+ "93": "hornbill",
114
+ "94": "hummingbird",
115
+ "95": "jacamar",
116
+ "96": "toucan",
117
+ "97": "drake",
118
+ "98": "red-breasted merganser, Mergus serrator",
119
+ "99": "goose",
120
+ "100": "black swan, Cygnus atratus",
121
+ "101": "tusker",
122
+ "102": "echidna, spiny anteater, anteater",
123
+ "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",
124
+ "104": "wallaby, brush kangaroo",
125
+ "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",
126
+ "106": "wombat",
127
+ "107": "jellyfish",
128
+ "108": "sea anemone, anemone",
129
+ "109": "brain coral",
130
+ "110": "flatworm, platyhelminth",
131
+ "111": "nematode, nematode worm, roundworm",
132
+ "112": "conch",
133
+ "113": "snail",
134
+ "114": "slug",
135
+ "115": "sea slug, nudibranch",
136
+ "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore",
137
+ "117": "chambered nautilus, pearly nautilus, nautilus",
138
+ "118": "Dungeness crab, Cancer magister",
139
+ "119": "rock crab, Cancer irroratus",
140
+ "120": "fiddler crab",
141
+ "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica",
142
+ "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus",
143
+ "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish",
144
+ "124": "crayfish, crawfish, crawdad, crawdaddy",
145
+ "125": "hermit crab",
146
+ "126": "isopod",
147
+ "127": "white stork, Ciconia ciconia",
148
+ "128": "black stork, Ciconia nigra",
149
+ "129": "spoonbill",
150
+ "130": "flamingo",
151
+ "131": "little blue heron, Egretta caerulea",
152
+ "132": "American egret, great white heron, Egretta albus",
153
+ "133": "bittern",
154
+ "134": "crane",
155
+ "135": "limpkin, Aramus pictus",
156
+ "136": "European gallinule, Porphyrio porphyrio",
157
+ "137": "American coot, marsh hen, mud hen, water hen, Fulica americana",
158
+ "138": "bustard",
159
+ "139": "ruddy turnstone, Arenaria interpres",
160
+ "140": "red-backed sandpiper, dunlin, Erolia alpina",
161
+ "141": "redshank, Tringa totanus",
162
+ "142": "dowitcher",
163
+ "143": "oystercatcher, oyster catcher",
164
+ "144": "pelican",
165
+ "145": "king penguin, Aptenodytes patagonica",
166
+ "146": "albatross, mollymawk",
167
+ "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus",
168
+ "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca",
169
+ "149": "dugong, Dugong dugon",
170
+ "150": "sea lion",
171
+ "151": "Chihuahua",
172
+ "152": "Japanese spaniel",
173
+ "153": "Maltese dog, Maltese terrier, Maltese",
174
+ "154": "Pekinese, Pekingese, Peke",
175
+ "155": "Shih-Tzu",
176
+ "156": "Blenheim spaniel",
177
+ "157": "papillon",
178
+ "158": "toy terrier",
179
+ "159": "Rhodesian ridgeback",
180
+ "160": "Afghan hound, Afghan",
181
+ "161": "basset, basset hound",
182
+ "162": "beagle",
183
+ "163": "bloodhound, sleuthhound",
184
+ "164": "bluetick",
185
+ "165": "black-and-tan coonhound",
186
+ "166": "Walker hound, Walker foxhound",
187
+ "167": "English foxhound",
188
+ "168": "redbone",
189
+ "169": "borzoi, Russian wolfhound",
190
+ "170": "Irish wolfhound",
191
+ "171": "Italian greyhound",
192
+ "172": "whippet",
193
+ "173": "Ibizan hound, Ibizan Podenco",
194
+ "174": "Norwegian elkhound, elkhound",
195
+ "175": "otterhound, otter hound",
196
+ "176": "Saluki, gazelle hound",
197
+ "177": "Scottish deerhound, deerhound",
198
+ "178": "Weimaraner",
199
+ "179": "Staffordshire bullterrier, Staffordshire bull terrier",
200
+ "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",
201
+ "181": "Bedlington terrier",
202
+ "182": "Border terrier",
203
+ "183": "Kerry blue terrier",
204
+ "184": "Irish terrier",
205
+ "185": "Norfolk terrier",
206
+ "186": "Norwich terrier",
207
+ "187": "Yorkshire terrier",
208
+ "188": "wire-haired fox terrier",
209
+ "189": "Lakeland terrier",
210
+ "190": "Sealyham terrier, Sealyham",
211
+ "191": "Airedale, Airedale terrier",
212
+ "192": "cairn, cairn terrier",
213
+ "193": "Australian terrier",
214
+ "194": "Dandie Dinmont, Dandie Dinmont terrier",
215
+ "195": "Boston bull, Boston terrier",
216
+ "196": "miniature schnauzer",
217
+ "197": "giant schnauzer",
218
+ "198": "standard schnauzer",
219
+ "199": "Scotch terrier, Scottish terrier, Scottie",
220
+ "200": "Tibetan terrier, chrysanthemum dog",
221
+ "201": "silky terrier, Sydney silky",
222
+ "202": "soft-coated wheaten terrier",
223
+ "203": "West Highland white terrier",
224
+ "204": "Lhasa, Lhasa apso",
225
+ "205": "flat-coated retriever",
226
+ "206": "curly-coated retriever",
227
+ "207": "golden retriever",
228
+ "208": "Labrador retriever",
229
+ "209": "Chesapeake Bay retriever",
230
+ "210": "German short-haired pointer",
231
+ "211": "vizsla, Hungarian pointer",
232
+ "212": "English setter",
233
+ "213": "Irish setter, red setter",
234
+ "214": "Gordon setter",
235
+ "215": "Brittany spaniel",
236
+ "216": "clumber, clumber spaniel",
237
+ "217": "English springer, English springer spaniel",
238
+ "218": "Welsh springer spaniel",
239
+ "219": "cocker spaniel, English cocker spaniel, cocker",
240
+ "220": "Sussex spaniel",
241
+ "221": "Irish water spaniel",
242
+ "222": "kuvasz",
243
+ "223": "schipperke",
244
+ "224": "groenendael",
245
+ "225": "malinois",
246
+ "226": "briard",
247
+ "227": "kelpie",
248
+ "228": "komondor",
249
+ "229": "Old English sheepdog, bobtail",
250
+ "230": "Shetland sheepdog, Shetland sheep dog, Shetland",
251
+ "231": "collie",
252
+ "232": "Border collie",
253
+ "233": "Bouvier des Flandres, Bouviers des Flandres",
254
+ "234": "Rottweiler",
255
+ "235": "German shepherd, German shepherd dog, German police dog, alsatian",
256
+ "236": "Doberman, Doberman pinscher",
257
+ "237": "miniature pinscher",
258
+ "238": "Greater Swiss Mountain dog",
259
+ "239": "Bernese mountain dog",
260
+ "240": "Appenzeller",
261
+ "241": "EntleBucher",
262
+ "242": "boxer",
263
+ "243": "bull mastiff",
264
+ "244": "Tibetan mastiff",
265
+ "245": "French bulldog",
266
+ "246": "Great Dane",
267
+ "247": "Saint Bernard, St Bernard",
268
+ "248": "Eskimo dog, husky",
269
+ "249": "malamute, malemute, Alaskan malamute",
270
+ "250": "Siberian husky",
271
+ "251": "dalmatian, coach dog, carriage dog",
272
+ "252": "affenpinscher, monkey pinscher, monkey dog",
273
+ "253": "basenji",
274
+ "254": "pug, pug-dog",
275
+ "255": "Leonberg",
276
+ "256": "Newfoundland, Newfoundland dog",
277
+ "257": "Great Pyrenees",
278
+ "258": "Samoyed, Samoyede",
279
+ "259": "Pomeranian",
280
+ "260": "chow, chow chow",
281
+ "261": "keeshond",
282
+ "262": "Brabancon griffon",
283
+ "263": "Pembroke, Pembroke Welsh corgi",
284
+ "264": "Cardigan, Cardigan Welsh corgi",
285
+ "265": "toy poodle",
286
+ "266": "miniature poodle",
287
+ "267": "standard poodle",
288
+ "268": "Mexican hairless",
289
+ "269": "timber wolf, grey wolf, gray wolf, Canis lupus",
290
+ "270": "white wolf, Arctic wolf, Canis lupus tundrarum",
291
+ "271": "red wolf, maned wolf, Canis rufus, Canis niger",
292
+ "272": "coyote, prairie wolf, brush wolf, Canis latrans",
293
+ "273": "dingo, warrigal, warragal, Canis dingo",
294
+ "274": "dhole, Cuon alpinus",
295
+ "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
296
+ "276": "hyena, hyaena",
297
+ "277": "red fox, Vulpes vulpes",
298
+ "278": "kit fox, Vulpes macrotis",
299
+ "279": "Arctic fox, white fox, Alopex lagopus",
300
+ "280": "grey fox, gray fox, Urocyon cinereoargenteus",
301
+ "281": "tabby, tabby cat",
302
+ "282": "tiger cat",
303
+ "283": "Persian cat",
304
+ "284": "Siamese cat, Siamese",
305
+ "285": "Egyptian cat",
306
+ "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor",
307
+ "287": "lynx, catamount",
308
+ "288": "leopard, Panthera pardus",
309
+ "289": "snow leopard, ounce, Panthera uncia",
310
+ "290": "jaguar, panther, Panthera onca, Felis onca",
311
+ "291": "lion, king of beasts, Panthera leo",
312
+ "292": "tiger, Panthera tigris",
313
+ "293": "cheetah, chetah, Acinonyx jubatus",
314
+ "294": "brown bear, bruin, Ursus arctos",
315
+ "295": "American black bear, black bear, Ursus americanus, Euarctos americanus",
316
+ "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus",
317
+ "297": "sloth bear, Melursus ursinus, Ursus ursinus",
318
+ "298": "mongoose",
319
+ "299": "meerkat, mierkat",
320
+ "300": "tiger beetle",
321
+ "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle",
322
+ "302": "ground beetle, carabid beetle",
323
+ "303": "long-horned beetle, longicorn, longicorn beetle",
324
+ "304": "leaf beetle, chrysomelid",
325
+ "305": "dung beetle",
326
+ "306": "rhinoceros beetle",
327
+ "307": "weevil",
328
+ "308": "fly",
329
+ "309": "bee",
330
+ "310": "ant, emmet, pismire",
331
+ "311": "grasshopper, hopper",
332
+ "312": "cricket",
333
+ "313": "walking stick, walkingstick, stick insect",
334
+ "314": "cockroach, roach",
335
+ "315": "mantis, mantid",
336
+ "316": "cicada, cicala",
337
+ "317": "leafhopper",
338
+ "318": "lacewing, lacewing fly",
339
+ "319": "dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",
340
+ "320": "damselfly",
341
+ "321": "admiral",
342
+ "322": "ringlet, ringlet butterfly",
343
+ "323": "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus",
344
+ "324": "cabbage butterfly",
345
+ "325": "sulphur butterfly, sulfur butterfly",
346
+ "326": "lycaenid, lycaenid butterfly",
347
+ "327": "starfish, sea star",
348
+ "328": "sea urchin",
349
+ "329": "sea cucumber, holothurian",
350
+ "330": "wood rabbit, cottontail, cottontail rabbit",
351
+ "331": "hare",
352
+ "332": "Angora, Angora rabbit",
353
+ "333": "hamster",
354
+ "334": "porcupine, hedgehog",
355
+ "335": "fox squirrel, eastern fox squirrel, Sciurus niger",
356
+ "336": "marmot",
357
+ "337": "beaver",
358
+ "338": "guinea pig, Cavia cobaya",
359
+ "339": "sorrel",
360
+ "340": "zebra",
361
+ "341": "hog, pig, grunter, squealer, Sus scrofa",
362
+ "342": "wild boar, boar, Sus scrofa",
363
+ "343": "warthog",
364
+ "344": "hippopotamus, hippo, river horse, Hippopotamus amphibius",
365
+ "345": "ox",
366
+ "346": "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis",
367
+ "347": "bison",
368
+ "348": "ram, tup",
369
+ "349": "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis",
370
+ "350": "ibex, Capra ibex",
371
+ "351": "hartebeest",
372
+ "352": "impala, Aepyceros melampus",
373
+ "353": "gazelle",
374
+ "354": "Arabian camel, dromedary, Camelus dromedarius",
375
+ "355": "llama",
376
+ "356": "weasel",
377
+ "357": "mink",
378
+ "358": "polecat, fitch, foulmart, foumart, Mustela putorius",
379
+ "359": "black-footed ferret, ferret, Mustela nigripes",
380
+ "360": "otter",
381
+ "361": "skunk, polecat, wood pussy",
382
+ "362": "badger",
383
+ "363": "armadillo",
384
+ "364": "three-toed sloth, ai, Bradypus tridactylus",
385
+ "365": "orangutan, orang, orangutang, Pongo pygmaeus",
386
+ "366": "gorilla, Gorilla gorilla",
387
+ "367": "chimpanzee, chimp, Pan troglodytes",
388
+ "368": "gibbon, Hylobates lar",
389
+ "369": "siamang, Hylobates syndactylus, Symphalangus syndactylus",
390
+ "370": "guenon, guenon monkey",
391
+ "371": "patas, hussar monkey, Erythrocebus patas",
392
+ "372": "baboon",
393
+ "373": "macaque",
394
+ "374": "langur",
395
+ "375": "colobus, colobus monkey",
396
+ "376": "proboscis monkey, Nasalis larvatus",
397
+ "377": "marmoset",
398
+ "378": "capuchin, ringtail, Cebus capucinus",
399
+ "379": "howler monkey, howler",
400
+ "380": "titi, titi monkey",
401
+ "381": "spider monkey, Ateles geoffroyi",
402
+ "382": "squirrel monkey, Saimiri sciureus",
403
+ "383": "Madagascar cat, ring-tailed lemur, Lemur catta",
404
+ "384": "indri, indris, Indri indri, Indri brevicaudatus",
405
+ "385": "Indian elephant, Elephas maximus",
406
+ "386": "African elephant, Loxodonta africana",
407
+ "387": "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens",
408
+ "388": "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca",
409
+ "389": "barracouta, snoek",
410
+ "390": "eel",
411
+ "391": "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch",
412
+ "392": "rock beauty, Holocanthus tricolor",
413
+ "393": "anemone fish",
414
+ "394": "sturgeon",
415
+ "395": "gar, garfish, garpike, billfish, Lepisosteus osseus",
416
+ "396": "lionfish",
417
+ "397": "puffer, pufferfish, blowfish, globefish",
418
+ "398": "abacus",
419
+ "399": "abaya",
420
+ "400": "academic gown, academic robe, judge's robe",
421
+ "401": "accordion, piano accordion, squeeze box",
422
+ "402": "acoustic guitar",
423
+ "403": "aircraft carrier, carrier, flattop, attack aircraft carrier",
424
+ "404": "airliner",
425
+ "405": "airship, dirigible",
426
+ "406": "altar",
427
+ "407": "ambulance",
428
+ "408": "amphibian, amphibious vehicle",
429
+ "409": "analog clock",
430
+ "410": "apiary, bee house",
431
+ "411": "apron",
432
+ "412": "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin",
433
+ "413": "assault rifle, assault gun",
434
+ "414": "backpack, back pack, knapsack, packsack, rucksack, haversack",
435
+ "415": "bakery, bakeshop, bakehouse",
436
+ "416": "balance beam, beam",
437
+ "417": "balloon",
438
+ "418": "ballpoint, ballpoint pen, ballpen, Biro",
439
+ "419": "Band Aid",
440
+ "420": "banjo",
441
+ "421": "bannister, banister, balustrade, balusters, handrail",
442
+ "422": "barbell",
443
+ "423": "barber chair",
444
+ "424": "barbershop",
445
+ "425": "barn",
446
+ "426": "barometer",
447
+ "427": "barrel, cask",
448
+ "428": "barrow, garden cart, lawn cart, wheelbarrow",
449
+ "429": "baseball",
450
+ "430": "basketball",
451
+ "431": "bassinet",
452
+ "432": "bassoon",
453
+ "433": "bathing cap, swimming cap",
454
+ "434": "bath towel",
455
+ "435": "bathtub, bathing tub, bath, tub",
456
+ "436": "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon",
457
+ "437": "beacon, lighthouse, beacon light, pharos",
458
+ "438": "beaker",
459
+ "439": "bearskin, busby, shako",
460
+ "440": "beer bottle",
461
+ "441": "beer glass",
462
+ "442": "bell cote, bell cot",
463
+ "443": "bib",
464
+ "444": "bicycle-built-for-two, tandem bicycle, tandem",
465
+ "445": "bikini, two-piece",
466
+ "446": "binder, ring-binder",
467
+ "447": "binoculars, field glasses, opera glasses",
468
+ "448": "birdhouse",
469
+ "449": "boathouse",
470
+ "450": "bobsled, bobsleigh, bob",
471
+ "451": "bolo tie, bolo, bola tie, bola",
472
+ "452": "bonnet, poke bonnet",
473
+ "453": "bookcase",
474
+ "454": "bookshop, bookstore, bookstall",
475
+ "455": "bottlecap",
476
+ "456": "bow",
477
+ "457": "bow tie, bow-tie, bowtie",
478
+ "458": "brass, memorial tablet, plaque",
479
+ "459": "brassiere, bra, bandeau",
480
+ "460": "breakwater, groin, groyne, mole, bulwark, seawall, jetty",
481
+ "461": "breastplate, aegis, egis",
482
+ "462": "broom",
483
+ "463": "bucket, pail",
484
+ "464": "buckle",
485
+ "465": "bulletproof vest",
486
+ "466": "bullet train, bullet",
487
+ "467": "butcher shop, meat market",
488
+ "468": "cab, hack, taxi, taxicab",
489
+ "469": "caldron, cauldron",
490
+ "470": "candle, taper, wax light",
491
+ "471": "cannon",
492
+ "472": "canoe",
493
+ "473": "can opener, tin opener",
494
+ "474": "cardigan",
495
+ "475": "car mirror",
496
+ "476": "carousel, carrousel, merry-go-round, roundabout, whirligig",
497
+ "477": "carpenter's kit, tool kit",
498
+ "478": "carton",
499
+ "479": "car wheel",
500
+ "480": "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM",
501
+ "481": "cassette",
502
+ "482": "cassette player",
503
+ "483": "castle",
504
+ "484": "catamaran",
505
+ "485": "CD player",
506
+ "486": "cello, violoncello",
507
+ "487": "cellular telephone, cellular phone, cellphone, cell, mobile phone",
508
+ "488": "chain",
509
+ "489": "chainlink fence",
510
+ "490": "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour",
511
+ "491": "chain saw, chainsaw",
512
+ "492": "chest",
513
+ "493": "chiffonier, commode",
514
+ "494": "chime, bell, gong",
515
+ "495": "china cabinet, china closet",
516
+ "496": "Christmas stocking",
517
+ "497": "church, church building",
518
+ "498": "cinema, movie theater, movie theatre, movie house, picture palace",
519
+ "499": "cleaver, meat cleaver, chopper",
520
+ "500": "cliff dwelling",
521
+ "501": "cloak",
522
+ "502": "clog, geta, patten, sabot",
523
+ "503": "cocktail shaker",
524
+ "504": "coffee mug",
525
+ "505": "coffeepot",
526
+ "506": "coil, spiral, volute, whorl, helix",
527
+ "507": "combination lock",
528
+ "508": "computer keyboard, keypad",
529
+ "509": "confectionery, confectionary, candy store",
530
+ "510": "container ship, containership, container vessel",
531
+ "511": "convertible",
532
+ "512": "corkscrew, bottle screw",
533
+ "513": "cornet, horn, trumpet, trump",
534
+ "514": "cowboy boot",
535
+ "515": "cowboy hat, ten-gallon hat",
536
+ "516": "cradle",
537
+ "517": "crane",
538
+ "518": "crash helmet",
539
+ "519": "crate",
540
+ "520": "crib, cot",
541
+ "521": "Crock Pot",
542
+ "522": "croquet ball",
543
+ "523": "crutch",
544
+ "524": "cuirass",
545
+ "525": "dam, dike, dyke",
546
+ "526": "desk",
547
+ "527": "desktop computer",
548
+ "528": "dial telephone, dial phone",
549
+ "529": "diaper, nappy, napkin",
550
+ "530": "digital clock",
551
+ "531": "digital watch",
552
+ "532": "dining table, board",
553
+ "533": "dishrag, dishcloth",
554
+ "534": "dishwasher, dish washer, dishwashing machine",
555
+ "535": "disk brake, disc brake",
556
+ "536": "dock, dockage, docking facility",
557
+ "537": "dogsled, dog sled, dog sleigh",
558
+ "538": "dome",
559
+ "539": "doormat, welcome mat",
560
+ "540": "drilling platform, offshore rig",
561
+ "541": "drum, membranophone, tympan",
562
+ "542": "drumstick",
563
+ "543": "dumbbell",
564
+ "544": "Dutch oven",
565
+ "545": "electric fan, blower",
566
+ "546": "electric guitar",
567
+ "547": "electric locomotive",
568
+ "548": "entertainment center",
569
+ "549": "envelope",
570
+ "550": "espresso maker",
571
+ "551": "face powder",
572
+ "552": "feather boa, boa",
573
+ "553": "file, file cabinet, filing cabinet",
574
+ "554": "fireboat",
575
+ "555": "fire engine, fire truck",
576
+ "556": "fire screen, fireguard",
577
+ "557": "flagpole, flagstaff",
578
+ "558": "flute, transverse flute",
579
+ "559": "folding chair",
580
+ "560": "football helmet",
581
+ "561": "forklift",
582
+ "562": "fountain",
583
+ "563": "fountain pen",
584
+ "564": "four-poster",
585
+ "565": "freight car",
586
+ "566": "French horn, horn",
587
+ "567": "frying pan, frypan, skillet",
588
+ "568": "fur coat",
589
+ "569": "garbage truck, dustcart",
590
+ "570": "gasmask, respirator, gas helmet",
591
+ "571": "gas pump, gasoline pump, petrol pump, island dispenser",
592
+ "572": "goblet",
593
+ "573": "go-kart",
594
+ "574": "golf ball",
595
+ "575": "golfcart, golf cart",
596
+ "576": "gondola",
597
+ "577": "gong, tam-tam",
598
+ "578": "gown",
599
+ "579": "grand piano, grand",
600
+ "580": "greenhouse, nursery, glasshouse",
601
+ "581": "grille, radiator grille",
602
+ "582": "grocery store, grocery, food market, market",
603
+ "583": "guillotine",
604
+ "584": "hair slide",
605
+ "585": "hair spray",
606
+ "586": "half track",
607
+ "587": "hammer",
608
+ "588": "hamper",
609
+ "589": "hand blower, blow dryer, blow drier, hair dryer, hair drier",
610
+ "590": "hand-held computer, hand-held microcomputer",
611
+ "591": "handkerchief, hankie, hanky, hankey",
612
+ "592": "hard disc, hard disk, fixed disk",
613
+ "593": "harmonica, mouth organ, harp, mouth harp",
614
+ "594": "harp",
615
+ "595": "harvester, reaper",
616
+ "596": "hatchet",
617
+ "597": "holster",
618
+ "598": "home theater, home theatre",
619
+ "599": "honeycomb",
620
+ "600": "hook, claw",
621
+ "601": "hoopskirt, crinoline",
622
+ "602": "horizontal bar, high bar",
623
+ "603": "horse cart, horse-cart",
624
+ "604": "hourglass",
625
+ "605": "iPod",
626
+ "606": "iron, smoothing iron",
627
+ "607": "jack-o'-lantern",
628
+ "608": "jean, blue jean, denim",
629
+ "609": "jeep, landrover",
630
+ "610": "jersey, T-shirt, tee shirt",
631
+ "611": "jigsaw puzzle",
632
+ "612": "jinrikisha, ricksha, rickshaw",
633
+ "613": "joystick",
634
+ "614": "kimono",
635
+ "615": "knee pad",
636
+ "616": "knot",
637
+ "617": "lab coat, laboratory coat",
638
+ "618": "ladle",
639
+ "619": "lampshade, lamp shade",
640
+ "620": "laptop, laptop computer",
641
+ "621": "lawn mower, mower",
642
+ "622": "lens cap, lens cover",
643
+ "623": "letter opener, paper knife, paperknife",
644
+ "624": "library",
645
+ "625": "lifeboat",
646
+ "626": "lighter, light, igniter, ignitor",
647
+ "627": "limousine, limo",
648
+ "628": "liner, ocean liner",
649
+ "629": "lipstick, lip rouge",
650
+ "630": "Loafer",
651
+ "631": "lotion",
652
+ "632": "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system",
653
+ "633": "loupe, jeweler's loupe",
654
+ "634": "lumbermill, sawmill",
655
+ "635": "magnetic compass",
656
+ "636": "mailbag, postbag",
657
+ "637": "mailbox, letter box",
658
+ "638": "maillot",
659
+ "639": "maillot, tank suit",
660
+ "640": "manhole cover",
661
+ "641": "maraca",
662
+ "642": "marimba, xylophone",
663
+ "643": "mask",
664
+ "644": "matchstick",
665
+ "645": "maypole",
666
+ "646": "maze, labyrinth",
667
+ "647": "measuring cup",
668
+ "648": "medicine chest, medicine cabinet",
669
+ "649": "megalith, megalithic structure",
670
+ "650": "microphone, mike",
671
+ "651": "microwave, microwave oven",
672
+ "652": "military uniform",
673
+ "653": "milk can",
674
+ "654": "minibus",
675
+ "655": "miniskirt, mini",
676
+ "656": "minivan",
677
+ "657": "missile",
678
+ "658": "mitten",
679
+ "659": "mixing bowl",
680
+ "660": "mobile home, manufactured home",
681
+ "661": "Model T",
682
+ "662": "modem",
683
+ "663": "monastery",
684
+ "664": "monitor",
685
+ "665": "moped",
686
+ "666": "mortar",
687
+ "667": "mortarboard",
688
+ "668": "mosque",
689
+ "669": "mosquito net",
690
+ "670": "motor scooter, scooter",
691
+ "671": "mountain bike, all-terrain bike, off-roader",
692
+ "672": "mountain tent",
693
+ "673": "mouse, computer mouse",
694
+ "674": "mousetrap",
695
+ "675": "moving van",
696
+ "676": "muzzle",
697
+ "677": "nail",
698
+ "678": "neck brace",
699
+ "679": "necklace",
700
+ "680": "nipple",
701
+ "681": "notebook, notebook computer",
702
+ "682": "obelisk",
703
+ "683": "oboe, hautboy, hautbois",
704
+ "684": "ocarina, sweet potato",
705
+ "685": "odometer, hodometer, mileometer, milometer",
706
+ "686": "oil filter",
707
+ "687": "organ, pipe organ",
708
+ "688": "oscilloscope, scope, cathode-ray oscilloscope, CRO",
709
+ "689": "overskirt",
710
+ "690": "oxcart",
711
+ "691": "oxygen mask",
712
+ "692": "packet",
713
+ "693": "paddle, boat paddle",
714
+ "694": "paddlewheel, paddle wheel",
715
+ "695": "padlock",
716
+ "696": "paintbrush",
717
+ "697": "pajama, pyjama, pj's, jammies",
718
+ "698": "palace",
719
+ "699": "panpipe, pandean pipe, syrinx",
720
+ "700": "paper towel",
721
+ "701": "parachute, chute",
722
+ "702": "parallel bars, bars",
723
+ "703": "park bench",
724
+ "704": "parking meter",
725
+ "705": "passenger car, coach, carriage",
726
+ "706": "patio, terrace",
727
+ "707": "pay-phone, pay-station",
728
+ "708": "pedestal, plinth, footstall",
729
+ "709": "pencil box, pencil case",
730
+ "710": "pencil sharpener",
731
+ "711": "perfume, essence",
732
+ "712": "Petri dish",
733
+ "713": "photocopier",
734
+ "714": "pick, plectrum, plectron",
735
+ "715": "pickelhaube",
736
+ "716": "picket fence, paling",
737
+ "717": "pickup, pickup truck",
738
+ "718": "pier",
739
+ "719": "piggy bank, penny bank",
740
+ "720": "pill bottle",
741
+ "721": "pillow",
742
+ "722": "ping-pong ball",
743
+ "723": "pinwheel",
744
+ "724": "pirate, pirate ship",
745
+ "725": "pitcher, ewer",
746
+ "726": "plane, carpenter's plane, woodworking plane",
747
+ "727": "planetarium",
748
+ "728": "plastic bag",
749
+ "729": "plate rack",
750
+ "730": "plow, plough",
751
+ "731": "plunger, plumber's helper",
752
+ "732": "Polaroid camera, Polaroid Land camera",
753
+ "733": "pole",
754
+ "734": "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria",
755
+ "735": "poncho",
756
+ "736": "pool table, billiard table, snooker table",
757
+ "737": "pop bottle, soda bottle",
758
+ "738": "pot, flowerpot",
759
+ "739": "potter's wheel",
760
+ "740": "power drill",
761
+ "741": "prayer rug, prayer mat",
762
+ "742": "printer",
763
+ "743": "prison, prison house",
764
+ "744": "projectile, missile",
765
+ "745": "projector",
766
+ "746": "puck, hockey puck",
767
+ "747": "punching bag, punch bag, punching ball, punchball",
768
+ "748": "purse",
769
+ "749": "quill, quill pen",
770
+ "750": "quilt, comforter, comfort, puff",
771
+ "751": "racer, race car, racing car",
772
+ "752": "racket, racquet",
773
+ "753": "radiator",
774
+ "754": "radio, wireless",
775
+ "755": "radio telescope, radio reflector",
776
+ "756": "rain barrel",
777
+ "757": "recreational vehicle, RV, R.V.",
778
+ "758": "reel",
779
+ "759": "reflex camera",
780
+ "760": "refrigerator, icebox",
781
+ "761": "remote control, remote",
782
+ "762": "restaurant, eating house, eating place, eatery",
783
+ "763": "revolver, six-gun, six-shooter",
784
+ "764": "rifle",
785
+ "765": "rocking chair, rocker",
786
+ "766": "rotisserie",
787
+ "767": "rubber eraser, rubber, pencil eraser",
788
+ "768": "rugby ball",
789
+ "769": "rule, ruler",
790
+ "770": "running shoe",
791
+ "771": "safe",
792
+ "772": "safety pin",
793
+ "773": "saltshaker, salt shaker",
794
+ "774": "sandal",
795
+ "775": "sarong",
796
+ "776": "sax, saxophone",
797
+ "777": "scabbard",
798
+ "778": "scale, weighing machine",
799
+ "779": "school bus",
800
+ "780": "schooner",
801
+ "781": "scoreboard",
802
+ "782": "screen, CRT screen",
803
+ "783": "screw",
804
+ "784": "screwdriver",
805
+ "785": "seat belt, seatbelt",
806
+ "786": "sewing machine",
807
+ "787": "shield, buckler",
808
+ "788": "shoe shop, shoe-shop, shoe store",
809
+ "789": "shoji",
810
+ "790": "shopping basket",
811
+ "791": "shopping cart",
812
+ "792": "shovel",
813
+ "793": "shower cap",
814
+ "794": "shower curtain",
815
+ "795": "ski",
816
+ "796": "ski mask",
817
+ "797": "sleeping bag",
818
+ "798": "slide rule, slipstick",
819
+ "799": "sliding door",
820
+ "800": "slot, one-armed bandit",
821
+ "801": "snorkel",
822
+ "802": "snowmobile",
823
+ "803": "snowplow, snowplough",
824
+ "804": "soap dispenser",
825
+ "805": "soccer ball",
826
+ "806": "sock",
827
+ "807": "solar dish, solar collector, solar furnace",
828
+ "808": "sombrero",
829
+ "809": "soup bowl",
830
+ "810": "space bar",
831
+ "811": "space heater",
832
+ "812": "space shuttle",
833
+ "813": "spatula",
834
+ "814": "speedboat",
835
+ "815": "spider web, spider's web",
836
+ "816": "spindle",
837
+ "817": "sports car, sport car",
838
+ "818": "spotlight, spot",
839
+ "819": "stage",
840
+ "820": "steam locomotive",
841
+ "821": "steel arch bridge",
842
+ "822": "steel drum",
843
+ "823": "stethoscope",
844
+ "824": "stole",
845
+ "825": "stone wall",
846
+ "826": "stopwatch, stop watch",
847
+ "827": "stove",
848
+ "828": "strainer",
849
+ "829": "streetcar, tram, tramcar, trolley, trolley car",
850
+ "830": "stretcher",
851
+ "831": "studio couch, day bed",
852
+ "832": "stupa, tope",
853
+ "833": "submarine, pigboat, sub, U-boat",
854
+ "834": "suit, suit of clothes",
855
+ "835": "sundial",
856
+ "836": "sunglass",
857
+ "837": "sunglasses, dark glasses, shades",
858
+ "838": "sunscreen, sunblock, sun blocker",
859
+ "839": "suspension bridge",
860
+ "840": "swab, swob, mop",
861
+ "841": "sweatshirt",
862
+ "842": "swimming trunks, bathing trunks",
863
+ "843": "swing",
864
+ "844": "switch, electric switch, electrical switch",
865
+ "845": "syringe",
866
+ "846": "table lamp",
867
+ "847": "tank, army tank, armored combat vehicle, armoured combat vehicle",
868
+ "848": "tape player",
869
+ "849": "teapot",
870
+ "850": "teddy, teddy bear",
871
+ "851": "television, television system",
872
+ "852": "tennis ball",
873
+ "853": "thatch, thatched roof",
874
+ "854": "theater curtain, theatre curtain",
875
+ "855": "thimble",
876
+ "856": "thresher, thrasher, threshing machine",
877
+ "857": "throne",
878
+ "858": "tile roof",
879
+ "859": "toaster",
880
+ "860": "tobacco shop, tobacconist shop, tobacconist",
881
+ "861": "toilet seat",
882
+ "862": "torch",
883
+ "863": "totem pole",
884
+ "864": "tow truck, tow car, wrecker",
885
+ "865": "toyshop",
886
+ "866": "tractor",
887
+ "867": "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi",
888
+ "868": "tray",
889
+ "869": "trench coat",
890
+ "870": "tricycle, trike, velocipede",
891
+ "871": "trimaran",
892
+ "872": "tripod",
893
+ "873": "triumphal arch",
894
+ "874": "trolleybus, trolley coach, trackless trolley",
895
+ "875": "trombone",
896
+ "876": "tub, vat",
897
+ "877": "turnstile",
898
+ "878": "typewriter keyboard",
899
+ "879": "umbrella",
900
+ "880": "unicycle, monocycle",
901
+ "881": "upright, upright piano",
902
+ "882": "vacuum, vacuum cleaner",
903
+ "883": "vase",
904
+ "884": "vault",
905
+ "885": "velvet",
906
+ "886": "vending machine",
907
+ "887": "vestment",
908
+ "888": "viaduct",
909
+ "889": "violin, fiddle",
910
+ "890": "volleyball",
911
+ "891": "waffle iron",
912
+ "892": "wall clock",
913
+ "893": "wallet, billfold, notecase, pocketbook",
914
+ "894": "wardrobe, closet, press",
915
+ "895": "warplane, military plane",
916
+ "896": "washbasin, handbasin, washbowl, lavabo, wash-hand basin",
917
+ "897": "washer, automatic washer, washing machine",
918
+ "898": "water bottle",
919
+ "899": "water jug",
920
+ "900": "water tower",
921
+ "901": "whiskey jug",
922
+ "902": "whistle",
923
+ "903": "wig",
924
+ "904": "window screen",
925
+ "905": "window shade",
926
+ "906": "Windsor tie",
927
+ "907": "wine bottle",
928
+ "908": "wing",
929
+ "909": "wok",
930
+ "910": "wooden spoon",
931
+ "911": "wool, woolen, woollen",
932
+ "912": "worm fence, snake fence, snake-rail fence, Virginia fence",
933
+ "913": "wreck",
934
+ "914": "yawl",
935
+ "915": "yurt",
936
+ "916": "web site, website, internet site, site",
937
+ "917": "comic book",
938
+ "918": "crossword puzzle, crossword",
939
+ "919": "street sign",
940
+ "920": "traffic light, traffic signal, stoplight",
941
+ "921": "book jacket, dust cover, dust jacket, dust wrapper",
942
+ "922": "menu",
943
+ "923": "plate",
944
+ "924": "guacamole",
945
+ "925": "consomme",
946
+ "926": "hot pot, hotpot",
947
+ "927": "trifle",
948
+ "928": "ice cream, icecream",
949
+ "929": "ice lolly, lolly, lollipop, popsicle",
950
+ "930": "French loaf",
951
+ "931": "bagel, beigel",
952
+ "932": "pretzel",
953
+ "933": "cheeseburger",
954
+ "934": "hotdog, hot dog, red hot",
955
+ "935": "mashed potato",
956
+ "936": "head cabbage",
957
+ "937": "broccoli",
958
+ "938": "cauliflower",
959
+ "939": "zucchini, courgette",
960
+ "940": "spaghetti squash",
961
+ "941": "acorn squash",
962
+ "942": "butternut squash",
963
+ "943": "cucumber, cuke",
964
+ "944": "artichoke, globe artichoke",
965
+ "945": "bell pepper",
966
+ "946": "cardoon",
967
+ "947": "mushroom",
968
+ "948": "Granny Smith",
969
+ "949": "strawberry",
970
+ "950": "orange",
971
+ "951": "lemon",
972
+ "952": "fig",
973
+ "953": "pineapple, ananas",
974
+ "954": "banana",
975
+ "955": "jackfruit, jak, jack",
976
+ "956": "custard apple",
977
+ "957": "pomegranate",
978
+ "958": "hay",
979
+ "959": "carbonara",
980
+ "960": "chocolate sauce, chocolate syrup",
981
+ "961": "dough",
982
+ "962": "meat loaf, meatloaf",
983
+ "963": "pizza, pizza pie",
984
+ "964": "potpie",
985
+ "965": "burrito",
986
+ "966": "red wine",
987
+ "967": "espresso",
988
+ "968": "cup",
989
+ "969": "eggnog",
990
+ "970": "alp",
991
+ "971": "bubble",
992
+ "972": "cliff, drop, drop-off",
993
+ "973": "coral reef",
994
+ "974": "geyser",
995
+ "975": "lakeside, lakeshore",
996
+ "976": "promontory, headland, head, foreland",
997
+ "977": "sandbar, sand bar",
998
+ "978": "seashore, coast, seacoast, sea-coast",
999
+ "979": "valley, vale",
1000
+ "980": "volcano",
1001
+ "981": "ballplayer, baseball player",
1002
+ "982": "groom, bridegroom",
1003
+ "983": "scuba diver",
1004
+ "984": "rapeseed",
1005
+ "985": "daisy",
1006
+ "986": "yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",
1007
+ "987": "corn",
1008
+ "988": "acorn",
1009
+ "989": "hip, rose hip, rosehip",
1010
+ "990": "buckeye, horse chestnut, conker",
1011
+ "991": "coral fungus",
1012
+ "992": "agaric",
1013
+ "993": "gyromitra",
1014
+ "994": "stinkhorn, carrion fungus",
1015
+ "995": "earthstar",
1016
+ "996": "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa",
1017
+ "997": "bolete",
1018
+ "998": "ear, spike, capitulum",
1019
+ "999": "toilet tissue, toilet paper, bathroom tissue"
1020
+ }
1021
+ }
DiT-MoE-B-8E2A/pipeline.py ADDED
@@ -0,0 +1,440 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import importlib
16
+ import inspect
17
+ import json
18
+ import sys
19
+ from pathlib import Path
20
+ from typing import Dict, List, Optional, Tuple, Union
21
+
22
+ import torch
23
+
24
+ from diffusers.image_processor import VaeImageProcessor
25
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
26
+ from diffusers.schedulers import DDIMScheduler, KarrasDiffusionSchedulers
27
+ from diffusers.utils import replace_example_docstring
28
+ from diffusers.utils.torch_utils import randn_tensor
29
+
30
+ DEFAULT_NATIVE_RESOLUTION = 256
31
+
32
+ EXAMPLE_DOC_STRING = """
33
+ Examples:
34
+ ```py
35
+ >>> from pathlib import Path
36
+ >>> import torch
37
+ >>> from diffusers import DiffusionPipeline
38
+
39
+ >>> model_dir = Path("./DiT-MoE-S-8E2A").resolve()
40
+ >>> pipe = DiffusionPipeline.from_pretrained(
41
+ ... str(model_dir),
42
+ ... local_files_only=True,
43
+ ... custom_pipeline=str(model_dir / "pipeline.py"),
44
+ ... trust_remote_code=True,
45
+ ... torch_dtype=torch.bfloat16,
46
+ ... )
47
+ >>> pipe.to("cuda")
48
+
49
+ >>> print(pipe.id2label[207])
50
+ >>> print(pipe.get_label_ids("golden retriever"))
51
+
52
+ >>> class_id = pipe.get_label_ids("golden retriever")[0]
53
+ >>> generator = torch.Generator(device="cuda").manual_seed(42)
54
+ >>> image = pipe(
55
+ ... class_labels=class_id,
56
+ ... height=256,
57
+ ... width=256,
58
+ ... num_inference_steps=50,
59
+ ... guidance_scale=4.0,
60
+ ... generator=generator,
61
+ ... ).images[0]
62
+ >>> image.save("demo.png")
63
+ ```
64
+ """
65
+
66
+
67
+ class DiTMoEPipeline(DiffusionPipeline):
68
+ r"""
69
+ Pipeline for class-conditional image generation with DiT-MoE.
70
+
71
+ Supports DDIM diffusion sampling and rectified-flow (RF) sampling.
72
+ Each checkpoint keeps an English `id2label` map in `model_index.json` (DiT-style).
73
+ """
74
+
75
+ model_cpu_offload_seq = "transformer->vae"
76
+ _optional_components = ["vae"]
77
+
78
+ def __init__(
79
+ self,
80
+ transformer,
81
+ scheduler: KarrasDiffusionSchedulers,
82
+ vae=None,
83
+ id2label: Optional[Dict[Union[int, str], str]] = None,
84
+ null_class_id: Optional[int] = None,
85
+ ):
86
+ super().__init__()
87
+ self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
88
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
89
+ self._use_rectified_flow = scheduler.__class__.__name__ == "DiTMoEFlowMatchScheduler"
90
+
91
+ if null_class_id is None:
92
+ null_class_id = int(getattr(self.transformer.config, "num_classes", 1000))
93
+ self.register_to_config(null_class_id=int(null_class_id))
94
+
95
+ self._id2label = self._normalize_id2label(id2label)
96
+ self.labels = self._build_label2id(self._id2label)
97
+
98
+ @property
99
+ def vae_scale_factor(self) -> int:
100
+ if self.vae is None:
101
+ return 8
102
+ block_out_channels = getattr(self.vae.config, "block_out_channels", None)
103
+ if block_out_channels:
104
+ return int(2 ** (len(block_out_channels) - 1))
105
+ return 8
106
+
107
+ @classmethod
108
+ def from_pretrained(cls, pretrained_model_name_or_path=None, subfolder=None, **kwargs):
109
+ r"""Load a self-contained variant folder locally or from the Hub."""
110
+ repo_root = Path(__file__).resolve().parent
111
+
112
+ if pretrained_model_name_or_path in (None, "", "."):
113
+ variant = repo_root
114
+ elif (
115
+ isinstance(pretrained_model_name_or_path, str)
116
+ and "/" in pretrained_model_name_or_path
117
+ and not Path(pretrained_model_name_or_path).exists()
118
+ ):
119
+ from huggingface_hub import snapshot_download
120
+
121
+ hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
122
+ if subfolder:
123
+ hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
124
+ cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
125
+ variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
126
+ else:
127
+ variant = Path(pretrained_model_name_or_path)
128
+ if not variant.is_absolute():
129
+ candidate = (Path.cwd() / variant).resolve()
130
+ variant = candidate if candidate.exists() else (repo_root / variant).resolve()
131
+ if subfolder:
132
+ variant = variant / subfolder
133
+
134
+ id2label_override = kwargs.pop("id2label", None)
135
+ null_class_id_override = kwargs.pop("null_class_id", None)
136
+ model_kwargs = dict(kwargs)
137
+ inserted: List[str] = []
138
+
139
+ def _load_component(folder: str, module_name: str, class_name: str):
140
+ comp_dir = variant / folder
141
+ module_path = comp_dir / f"{module_name}.py"
142
+ has_weights = (comp_dir / "config.json").exists() or (comp_dir / "scheduler_config.json").exists()
143
+ if not module_path.exists() or not has_weights:
144
+ return None
145
+
146
+ comp_path = str(comp_dir)
147
+ if comp_path not in sys.path:
148
+ sys.path.insert(0, comp_path)
149
+ inserted.append(comp_path)
150
+
151
+ module = importlib.import_module(module_name)
152
+ component_cls = getattr(module, class_name)
153
+ return component_cls.from_pretrained(str(comp_dir), **model_kwargs)
154
+
155
+ try:
156
+ transformer = _load_component("transformer", "transformer_dit_moe", "DiTMoETransformer2DModel")
157
+ if transformer is None:
158
+ raise ValueError(f"No loadable transformer found under {variant}")
159
+
160
+ scheduler = None
161
+ scheduler_config_path = variant / "scheduler" / "scheduler_config.json"
162
+ scheduler_dir = variant / "scheduler"
163
+ if scheduler_config_path.exists():
164
+ scheduler_config = json.loads(scheduler_config_path.read_text(encoding="utf-8"))
165
+ scheduler_class = scheduler_config.get("_class_name", "DDIMScheduler")
166
+ if scheduler_class == "DiTMoEFlowMatchScheduler":
167
+ scheduler_path = str(scheduler_dir)
168
+ if scheduler_path not in sys.path:
169
+ sys.path.insert(0, scheduler_path)
170
+ inserted.append(scheduler_path)
171
+ scheduler_module = importlib.import_module("scheduling_flow_match_dit_moe")
172
+ scheduler_cls = getattr(scheduler_module, "DiTMoEFlowMatchScheduler")
173
+ scheduler = scheduler_cls.from_pretrained(str(variant), subfolder="scheduler")
174
+ else:
175
+ scheduler = DDIMScheduler.from_pretrained(str(variant), subfolder="scheduler")
176
+ if scheduler is None:
177
+ scheduler = DDIMScheduler(num_train_timesteps=1000)
178
+
179
+ vae = None
180
+ vae_dir = variant / "vae"
181
+ if vae_dir.exists() and (vae_dir / "config.json").exists():
182
+ from diffusers import AutoencoderKL
183
+
184
+ vae = AutoencoderKL.from_pretrained(str(vae_dir), **model_kwargs)
185
+
186
+ pipeline_config = cls._read_pipeline_config_from_model_index(str(variant))
187
+ id2label = id2label_override or pipeline_config.get("id2label")
188
+ null_class_id = null_class_id_override if null_class_id_override is not None else pipeline_config.get(
189
+ "null_class_id"
190
+ )
191
+ pipe = cls(
192
+ transformer=transformer,
193
+ scheduler=scheduler,
194
+ vae=vae,
195
+ id2label=id2label,
196
+ null_class_id=null_class_id,
197
+ )
198
+ if hasattr(pipe, "register_to_config"):
199
+ pipe.register_to_config(_name_or_path=str(variant))
200
+ return pipe
201
+ finally:
202
+ for comp_path in inserted:
203
+ if comp_path in sys.path:
204
+ sys.path.remove(comp_path)
205
+
206
+ @staticmethod
207
+ def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
208
+ if not id2label:
209
+ return {}
210
+ return {int(key): value for key, value in id2label.items()}
211
+
212
+ @staticmethod
213
+ def _read_pipeline_config_from_model_index(variant_path: Optional[str]) -> Dict[str, object]:
214
+ if not variant_path:
215
+ return {}
216
+ variant_dir = Path(variant_path).resolve()
217
+ model_index_path = variant_dir / "model_index.json"
218
+ if not model_index_path.exists():
219
+ return {}
220
+ raw = json.loads(model_index_path.read_text(encoding="utf-8"))
221
+ config: Dict[str, object] = {}
222
+ id2label = raw.get("id2label")
223
+ if isinstance(id2label, dict):
224
+ config["id2label"] = {int(key): value for key, value in id2label.items()}
225
+ if "null_class_id" in raw:
226
+ config["null_class_id"] = int(raw["null_class_id"])
227
+ return config
228
+
229
+ @staticmethod
230
+ def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]:
231
+ label2id: Dict[str, int] = {}
232
+ for class_id, value in id2label.items():
233
+ for synonym in value.split(","):
234
+ synonym = synonym.strip()
235
+ if synonym:
236
+ label2id[synonym] = int(class_id)
237
+ return dict(sorted(label2id.items()))
238
+
239
+ @property
240
+ def id2label(self) -> Dict[int, str]:
241
+ return self._id2label
242
+
243
+ def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
244
+ r"""Map English ImageNet labels to class ids."""
245
+ labels = [label] if isinstance(label, str) else label
246
+ if not self.labels:
247
+ raise ValueError("No id2label mapping is available in this checkpoint.")
248
+ missing = [item for item in labels if item not in self.labels]
249
+ if missing:
250
+ preview = ", ".join(list(self.labels.keys())[:8])
251
+ raise ValueError(f"Unknown labels: {missing}. Example valid labels: {preview}, ...")
252
+ return [self.labels[item] for item in labels]
253
+
254
+ def _normalize_class_labels(
255
+ self,
256
+ class_labels: Union[int, str, List[Union[int, str]], torch.Tensor],
257
+ ) -> List[int]:
258
+ if isinstance(class_labels, torch.Tensor):
259
+ class_labels = class_labels.detach().cpu().tolist()
260
+ if isinstance(class_labels, int):
261
+ return [class_labels]
262
+ if isinstance(class_labels, str):
263
+ return self.get_label_ids(class_labels)
264
+ if not class_labels:
265
+ raise ValueError("`class_labels` cannot be empty.")
266
+ if isinstance(class_labels[0], str):
267
+ return self.get_label_ids(class_labels) # type: ignore[arg-type]
268
+ return [int(class_id) for class_id in class_labels] # type: ignore[union-attr]
269
+
270
+ def prepare_latents(
271
+ self,
272
+ batch_size: int,
273
+ num_channels: int,
274
+ height: int,
275
+ width: int,
276
+ dtype: torch.dtype,
277
+ device: torch.device,
278
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]],
279
+ ) -> torch.Tensor:
280
+ latent_height = height // self.vae_scale_factor
281
+ latent_width = width // self.vae_scale_factor
282
+ return randn_tensor(
283
+ (batch_size, num_channels, latent_height, latent_width),
284
+ generator=generator,
285
+ device=device,
286
+ dtype=dtype,
287
+ )
288
+
289
+ @staticmethod
290
+ def prepare_extra_step_kwargs(
291
+ scheduler: KarrasDiffusionSchedulers,
292
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]],
293
+ eta: float,
294
+ ) -> Dict[str, object]:
295
+ kwargs: Dict[str, object] = {}
296
+ step_params = set(inspect.signature(scheduler.step).parameters.keys())
297
+ if "eta" in step_params:
298
+ kwargs["eta"] = eta
299
+ if "generator" in step_params:
300
+ kwargs["generator"] = generator
301
+ return kwargs
302
+
303
+ def _apply_cfg(self, model_output: torch.Tensor, guidance_scale: float) -> torch.Tensor:
304
+ if guidance_scale <= 1.0:
305
+ return model_output
306
+ cond, uncond = model_output.chunk(2)
307
+ if self.transformer.learn_sigma:
308
+ eps_cond, rest_cond = cond[:, : self.transformer.in_channels], cond[:, self.transformer.in_channels :]
309
+ eps_uncond, _rest_uncond = uncond[:, : self.transformer.in_channels], uncond[:, self.transformer.in_channels :]
310
+ eps = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
311
+ guided = torch.cat([eps, rest_cond], dim=1)
312
+ else:
313
+ guided = uncond + guidance_scale * (cond - uncond)
314
+ return guided
315
+
316
+ def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"):
317
+ if self.vae is None:
318
+ if output_type == "latent":
319
+ return latents
320
+ raise ValueError("Cannot decode latents without a VAE.")
321
+ scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215)
322
+ latents = latents / scaling_factor
323
+ if output_type == "latent":
324
+ return latents
325
+ image = self.vae.decode(latents).sample
326
+ return self.image_processor.postprocess(image, output_type=output_type)
327
+
328
+ @torch.inference_mode()
329
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
330
+ def __call__(
331
+ self,
332
+ class_labels: Union[int, str, List[Union[int, str]], torch.Tensor] = 207,
333
+ height: Optional[int] = None,
334
+ width: Optional[int] = None,
335
+ num_inference_steps: int = 50,
336
+ guidance_scale: float = 4.0,
337
+ eta: float = 0.0,
338
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
339
+ output_type: str = "pil",
340
+ return_dict: bool = True,
341
+ ) -> Union[ImagePipelineOutput, Tuple]:
342
+ r"""
343
+ Generate class-conditional samples from a DiT-MoE checkpoint.
344
+
345
+ Examples:
346
+ <!-- this section is replaced by replace_example_docstring -->
347
+ """
348
+ class_labels_list = self._normalize_class_labels(class_labels)
349
+ batch_size = len(class_labels_list)
350
+ native_size = int(getattr(self.transformer.config, "input_size", 32)) * self.vae_scale_factor
351
+ height = native_size if height is None else int(height)
352
+ width = native_size if width is None else int(width)
353
+
354
+ if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
355
+ raise ValueError(
356
+ f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})."
357
+ )
358
+ if output_type not in {"pil", "np", "pt", "latent"}:
359
+ raise ValueError(f"Unsupported `output_type`: {output_type}")
360
+
361
+ device = self._execution_device
362
+ model_dtype = next(self.transformer.parameters()).dtype
363
+ latents = self.prepare_latents(
364
+ batch_size,
365
+ self.transformer.config.in_channels,
366
+ height,
367
+ width,
368
+ model_dtype,
369
+ device,
370
+ generator,
371
+ )
372
+
373
+ do_cfg = guidance_scale > 1.0
374
+ class_labels_tensor = torch.tensor(class_labels_list, device=device, dtype=torch.long)
375
+ if do_cfg:
376
+ latents = torch.cat([latents, latents], dim=0)
377
+ null_class = int(self.config.null_class_id)
378
+ uncond = torch.full((batch_size,), null_class, device=device, dtype=torch.long)
379
+ labels = torch.cat([class_labels_tensor, uncond], dim=0)
380
+ else:
381
+ labels = class_labels_tensor
382
+
383
+ extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator, eta=eta)
384
+
385
+ if self._use_rectified_flow:
386
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
387
+ timesteps = self.scheduler.timesteps
388
+ for index, timestep in enumerate(self.progress_bar(timesteps)):
389
+ next_timestep = timesteps[index + 1] if index + 1 < len(timesteps) else torch.tensor(
390
+ 0.0, device=device
391
+ )
392
+ timestep_batch = torch.full((labels.shape[0],), float(timestep), device=device, dtype=model_dtype)
393
+ model_output = self.transformer(latents, timestep_batch, labels, return_dict=True).sample
394
+ model_output = self.transformer.split_velocity(model_output)
395
+ if do_cfg:
396
+ model_output = self._apply_cfg(model_output, guidance_scale)
397
+ latents_cfg = latents[:batch_size]
398
+ else:
399
+ latents_cfg = latents
400
+ step_output = self.scheduler.step(
401
+ model_output[:batch_size] if do_cfg else model_output,
402
+ timestep_batch[:batch_size] if do_cfg else timestep_batch,
403
+ latents_cfg,
404
+ next_timestep=next_timestep,
405
+ ).prev_sample
406
+ latents = step_output if not do_cfg else torch.cat([step_output, step_output], dim=0)
407
+ latents = latents[:batch_size]
408
+ else:
409
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
410
+ for timestep in self.progress_bar(self.scheduler.timesteps):
411
+ timestep_batch = torch.full((labels.shape[0],), float(timestep), device=device, dtype=model_dtype)
412
+ latent_model_input = self.scheduler.scale_model_input(latents, timestep)
413
+ model_output = self.transformer(latent_model_input, timestep_batch, labels, return_dict=True).sample
414
+ if do_cfg:
415
+ model_output = self._apply_cfg(model_output, guidance_scale)
416
+ latents_input = latents[:batch_size]
417
+ model_output = model_output[:batch_size]
418
+ else:
419
+ latents_input = latents
420
+ if self.transformer.learn_sigma:
421
+ model_output, _ = torch.split(model_output, self.transformer.in_channels, dim=1)
422
+ latents = self.scheduler.step(
423
+ model_output,
424
+ timestep,
425
+ latents_input,
426
+ **extra_step_kwargs,
427
+ ).prev_sample
428
+ if do_cfg:
429
+ latents = torch.cat([latents, latents], dim=0)
430
+
431
+ image = self.decode_latents(latents, output_type=output_type)
432
+ self.maybe_free_model_hooks()
433
+ if not return_dict:
434
+ return (image,)
435
+ return ImagePipelineOutput(images=image)
436
+
437
+
438
+ DiTMoEPipelineOutput = ImagePipelineOutput
439
+
440
+ __all__ = ["DiTMoEPipeline", "DiTMoEPipelineOutput"]
DiT-MoE-B-8E2A/scheduler/scheduler_config.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "DDIMScheduler",
3
+ "_diffusers_version": "0.36.0",
4
+ "beta_end": 0.02,
5
+ "beta_schedule": "linear",
6
+ "beta_start": 0.0001,
7
+ "clip_sample": false,
8
+ "num_train_timesteps": 1000,
9
+ "prediction_type": "epsilon",
10
+ "set_alpha_to_one": true,
11
+ "steps_offset": 0,
12
+ "trained_betas": null
13
+ }
DiT-MoE-B-8E2A/transformer/config.json ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "DiTMoETransformer2DModel",
3
+ "class_dropout_prob": 0.1,
4
+ "depth": 12,
5
+ "hidden_size": 768,
6
+ "in_channels": 4,
7
+ "input_size": 32,
8
+ "learn_sigma": true,
9
+ "num_classes": 1000,
10
+ "num_experts": 8,
11
+ "num_experts_per_tok": 2,
12
+ "num_heads": 12,
13
+ "patch_size": 2,
14
+ "pretraining_tp": 2,
15
+ "use_flash_attn": false
16
+ }
DiT-MoE-B-8E2A/transformer/diffusion_pytorch_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b515c7400d23980bc73f178cde2d7b35bfbb1dc0e43ccf24284f981e7fac40af
3
+ size 3183493168
DiT-MoE-B-8E2A/transformer/transformer_dit_moe.py ADDED
@@ -0,0 +1,511 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+
7
+ from dataclasses import dataclass
8
+ import math
9
+ from typing import Optional, Tuple, Union
10
+
11
+ import numpy as np
12
+ import torch
13
+ import torch.nn as nn
14
+ import torch.nn.functional as F
15
+
16
+ try:
17
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
18
+ from diffusers.models.modeling_utils import ModelMixin
19
+ from diffusers.utils import BaseOutput
20
+ except Exception: # pragma: no cover
21
+ class BaseOutput(dict):
22
+ def __post_init__(self):
23
+ self.update(self.__dict__)
24
+
25
+ class _Config(dict):
26
+ def __getattr__(self, key):
27
+ try:
28
+ return self[key]
29
+ except KeyError as error:
30
+ raise AttributeError(key) from error
31
+
32
+ class ConfigMixin:
33
+ config_name = "config.json"
34
+
35
+ class ModelMixin(nn.Module):
36
+ pass
37
+
38
+ def register_to_config(init):
39
+ def wrapper(self, *args, **kwargs):
40
+ import inspect
41
+
42
+ signature = inspect.signature(init)
43
+ bound = signature.bind(self, *args, **kwargs)
44
+ bound.apply_defaults()
45
+ self.config = _Config({key: value for key, value in bound.arguments.items() if key != "self"})
46
+ init(self, *args, **kwargs)
47
+
48
+ return wrapper
49
+
50
+
51
+ try:
52
+ from timm.models.vision_transformer import Attention
53
+ except Exception: # pragma: no cover
54
+ Attention = None
55
+
56
+ try:
57
+ import flash_attn
58
+
59
+ if hasattr(flash_attn, "__version__") and int(flash_attn.__version__[0]) == 2:
60
+ from flash_attn.modules.mha import FlashSelfAttention
61
+ else:
62
+ from flash_attn.modules.mha import FlashSelfAttention
63
+ except Exception:
64
+ FlashSelfAttention = None
65
+
66
+
67
+ @dataclass
68
+ class DiTMoETransformer2DModelOutput(BaseOutput):
69
+ sample: torch.FloatTensor
70
+
71
+
72
+ def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
73
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
74
+
75
+
76
+ class DiTMoEPatchEmbed(nn.Module):
77
+ def __init__(self, input_size: int, patch_size: int, in_channels: int, embed_dim: int, bias: bool = True):
78
+ super().__init__()
79
+ self.input_size = (input_size, input_size)
80
+ self.patch_size = (patch_size, patch_size)
81
+ self.num_patches = (input_size // patch_size) ** 2
82
+ self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
83
+
84
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
85
+ hidden_states = self.proj(hidden_states)
86
+ return hidden_states.flatten(2).transpose(1, 2)
87
+
88
+
89
+ class DiTMoETimestepEmbedder(nn.Module):
90
+ def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
91
+ super().__init__()
92
+ self.mlp = nn.Sequential(
93
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
94
+ nn.SiLU(),
95
+ nn.Linear(hidden_size, hidden_size, bias=True),
96
+ )
97
+ self.frequency_embedding_size = frequency_embedding_size
98
+
99
+ @staticmethod
100
+ def timestep_embedding(timesteps: torch.Tensor, dim: int, max_period: int = 10000):
101
+ half = dim // 2
102
+ freqs = torch.exp(
103
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half
104
+ )
105
+ args = timesteps[:, None].float() * freqs[None]
106
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
107
+ if dim % 2:
108
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
109
+ return embedding
110
+
111
+ def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
112
+ t_freq = self.timestep_embedding(timesteps, self.frequency_embedding_size)
113
+ weight_dtype = self.mlp[0].weight.dtype
114
+ return self.mlp(t_freq.to(dtype=weight_dtype))
115
+
116
+
117
+ class DiTMoELabelEmbedder(nn.Module):
118
+ def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float):
119
+ super().__init__()
120
+ use_cfg_embedding = dropout_prob > 0
121
+ self.embedding_table = nn.Embedding(num_classes + int(use_cfg_embedding), hidden_size)
122
+ self.num_classes = num_classes
123
+ self.dropout_prob = dropout_prob
124
+
125
+ def token_drop(self, labels: torch.Tensor, force_drop_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
126
+ if force_drop_ids is None:
127
+ drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
128
+ else:
129
+ drop_ids = force_drop_ids == 1
130
+ return torch.where(drop_ids, self.num_classes, labels)
131
+
132
+ def forward(
133
+ self,
134
+ labels: torch.Tensor,
135
+ train: bool = False,
136
+ force_drop_ids: Optional[torch.Tensor] = None,
137
+ ) -> torch.Tensor:
138
+ use_dropout = self.dropout_prob > 0
139
+ if (train and use_dropout) or (force_drop_ids is not None):
140
+ labels = self.token_drop(labels, force_drop_ids)
141
+ return self.embedding_table(labels)
142
+
143
+
144
+ class MoEGate(nn.Module):
145
+ def __init__(self, embed_dim: int, num_experts: int = 16, num_experts_per_tok: int = 2, aux_loss_alpha: float = 0.01):
146
+ super().__init__()
147
+ self.top_k = num_experts_per_tok
148
+ self.n_routed_experts = num_experts
149
+ self.scoring_func = "softmax"
150
+ self.alpha = aux_loss_alpha
151
+ self.seq_aux = False
152
+ self.norm_topk_prob = False
153
+ self.gating_dim = embed_dim
154
+ self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
155
+ self.reset_parameters()
156
+
157
+ def reset_parameters(self) -> None:
158
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
159
+
160
+ def forward(self, hidden_states: torch.Tensor):
161
+ bsz, seq_len, hidden_dim = hidden_states.shape
162
+ hidden_states = hidden_states.view(-1, hidden_dim)
163
+ logits = F.linear(hidden_states, self.weight, None)
164
+ scores = logits.softmax(dim=-1)
165
+ topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
166
+ if self.top_k > 1 and self.norm_topk_prob:
167
+ topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
168
+
169
+ if self.training and self.alpha > 0.0:
170
+ topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
171
+ mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
172
+ ce = mask_ce.float().mean(0)
173
+ pi = scores.mean(0)
174
+ fi = ce * self.n_routed_experts
175
+ aux_loss = (pi * fi).sum() * self.alpha
176
+ else:
177
+ aux_loss = None
178
+ return topk_idx, topk_weight, aux_loss
179
+
180
+
181
+ class AddAuxiliaryLoss(torch.autograd.Function):
182
+ @staticmethod
183
+ def forward(ctx, x, loss):
184
+ assert loss.numel() == 1
185
+ ctx.dtype = loss.dtype
186
+ ctx.required_aux_loss = loss.requires_grad
187
+ return x
188
+
189
+ @staticmethod
190
+ def backward(ctx, grad_output):
191
+ grad_loss = None
192
+ if ctx.required_aux_loss:
193
+ grad_loss = torch.ones(1, dtype=ctx.dtype, device=grad_output.device)
194
+ return grad_output, grad_loss
195
+
196
+
197
+ class MoeMLP(nn.Module):
198
+ def __init__(self, hidden_size: int, intermediate_size: int, pretraining_tp: int = 2):
199
+ super().__init__()
200
+ self.hidden_size = hidden_size
201
+ self.intermediate_size = intermediate_size
202
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
203
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
204
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
205
+ self.act_fn = nn.SiLU()
206
+ self.pretraining_tp = pretraining_tp
207
+
208
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
209
+ if self.pretraining_tp > 1:
210
+ slice_size = self.intermediate_size // self.pretraining_tp
211
+ gate_proj_slices = self.gate_proj.weight.split(slice_size, dim=0)
212
+ up_proj_slices = self.up_proj.weight.split(slice_size, dim=0)
213
+ down_proj_slices = self.down_proj.weight.split(slice_size, dim=1)
214
+ gate_proj = torch.cat([F.linear(x, gate_proj_slices[i]) for i in range(self.pretraining_tp)], dim=-1)
215
+ up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.pretraining_tp)], dim=-1)
216
+ intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice_size, dim=-1)
217
+ down_proj = [F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.pretraining_tp)]
218
+ return sum(down_proj)
219
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
220
+
221
+
222
+ class SparseMoeBlock(nn.Module):
223
+ def __init__(
224
+ self,
225
+ embed_dim: int,
226
+ mlp_ratio: int = 4,
227
+ num_experts: int = 16,
228
+ num_experts_per_tok: int = 2,
229
+ pretraining_tp: int = 2,
230
+ ):
231
+ super().__init__()
232
+ self.num_experts_per_tok = num_experts_per_tok
233
+ self.experts = nn.ModuleList(
234
+ [
235
+ MoeMLP(
236
+ hidden_size=embed_dim,
237
+ intermediate_size=mlp_ratio * embed_dim,
238
+ pretraining_tp=pretraining_tp,
239
+ )
240
+ for _ in range(num_experts)
241
+ ]
242
+ )
243
+ self.gate = MoEGate(embed_dim=embed_dim, num_experts=num_experts, num_experts_per_tok=num_experts_per_tok)
244
+ self.n_shared_experts = 2
245
+ if self.n_shared_experts is not None:
246
+ intermediate_size = embed_dim * self.n_shared_experts
247
+ self.shared_experts = MoeMLP(
248
+ hidden_size=embed_dim,
249
+ intermediate_size=intermediate_size,
250
+ pretraining_tp=pretraining_tp,
251
+ )
252
+
253
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
254
+ identity = hidden_states
255
+ orig_shape = hidden_states.shape
256
+ topk_idx, topk_weight, aux_loss = self.gate(hidden_states)
257
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
258
+ flat_topk_idx = topk_idx.view(-1)
259
+ if self.training:
260
+ hidden_states = hidden_states.repeat_interleave(self.num_experts_per_tok, dim=0)
261
+ y = torch.empty_like(hidden_states, dtype=hidden_states.dtype)
262
+ for index, expert in enumerate(self.experts):
263
+ y[flat_topk_idx == index] = expert(hidden_states[flat_topk_idx == index]).float()
264
+ y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
265
+ y = y.view(*orig_shape)
266
+ y = AddAuxiliaryLoss.apply(y, aux_loss)
267
+ else:
268
+ y = self.moe_infer(hidden_states, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
269
+ if self.n_shared_experts is not None:
270
+ y = y + self.shared_experts(identity)
271
+ return y
272
+
273
+ @torch.no_grad()
274
+ def moe_infer(self, x: torch.Tensor, flat_expert_indices: torch.Tensor, flat_expert_weights: torch.Tensor):
275
+ expert_cache = torch.zeros_like(x)
276
+ idxs = flat_expert_indices.argsort()
277
+ tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
278
+ token_idxs = idxs // self.num_experts_per_tok
279
+ for index, end_idx in enumerate(tokens_per_expert):
280
+ start_idx = 0 if index == 0 else tokens_per_expert[index - 1]
281
+ if start_idx == end_idx:
282
+ continue
283
+ expert = self.experts[index]
284
+ exp_token_idx = token_idxs[start_idx:end_idx]
285
+ expert_tokens = x[exp_token_idx]
286
+ expert_out = expert(expert_tokens)
287
+ expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
288
+ expert_cache = expert_cache.to(expert_out.dtype)
289
+ expert_cache.scatter_reduce_(
290
+ 0,
291
+ exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]),
292
+ expert_out,
293
+ reduce="sum",
294
+ )
295
+ return expert_cache
296
+
297
+
298
+ class FlashSelfMHAModified(nn.Module):
299
+ def __init__(self, dim: int, num_heads: int, qkv_bias: bool = True, qk_norm: bool = False):
300
+ super().__init__()
301
+ if FlashSelfAttention is None:
302
+ raise ImportError("flash_attn is required when use_flash_attn=True.")
303
+ self.dim = dim
304
+ self.num_heads = num_heads
305
+ self.head_dim = dim // num_heads
306
+ self.Wqkv = nn.Linear(dim, 3 * dim, bias=qkv_bias)
307
+ self.q_norm = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
308
+ self.k_norm = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
309
+ self.inner_attn = FlashSelfAttention(attention_dropout=0.0)
310
+ self.out_proj = nn.Linear(dim, dim, bias=qkv_bias)
311
+ self.proj_drop = nn.Dropout(0.0)
312
+
313
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
314
+ batch_size, seq_len, dim = x.shape
315
+ qkv = self.Wqkv(x).view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
316
+ query, key, value = qkv.unbind(dim=2)
317
+ query = self.q_norm(query).to(dtype=torch.float16)
318
+ key = self.k_norm(key).to(dtype=torch.float16)
319
+ qkv = torch.stack([query, key, value], dim=2)
320
+ context = self.inner_attn(qkv)
321
+ out = self.out_proj(context.view(batch_size, seq_len, dim))
322
+ return self.proj_drop(out)
323
+
324
+
325
+ class DiTMoEBlock(nn.Module):
326
+ def __init__(
327
+ self,
328
+ hidden_size: int,
329
+ num_heads: int,
330
+ mlp_ratio: float = 4.0,
331
+ num_experts: int = 8,
332
+ num_experts_per_tok: int = 2,
333
+ pretraining_tp: int = 2,
334
+ use_flash_attn: bool = False,
335
+ ):
336
+ super().__init__()
337
+ self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
338
+ if use_flash_attn:
339
+ self.attn = FlashSelfMHAModified(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True)
340
+ else:
341
+ if Attention is None:
342
+ raise ImportError("timm is required when use_flash_attn=False.")
343
+ self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
344
+ self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
345
+ self.moe = SparseMoeBlock(hidden_size, int(mlp_ratio), num_experts, num_experts_per_tok, pretraining_tp)
346
+ self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
347
+
348
+ def forward(self, hidden_states: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
349
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(conditioning).chunk(
350
+ 6, dim=1
351
+ )
352
+ hidden_states = hidden_states + gate_msa.unsqueeze(1) * self.attn(
353
+ modulate(self.norm1(hidden_states), shift_msa, scale_msa)
354
+ )
355
+ hidden_states = hidden_states + gate_mlp.unsqueeze(1) * self.moe(
356
+ modulate(self.norm2(hidden_states), shift_mlp, scale_mlp)
357
+ )
358
+ return hidden_states
359
+
360
+
361
+ class DiTMoEFinalLayer(nn.Module):
362
+ def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
363
+ super().__init__()
364
+ self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
365
+ self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
366
+ self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
367
+
368
+ def forward(self, hidden_states: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
369
+ shift, scale = self.adaLN_modulation(conditioning).chunk(2, dim=1)
370
+ hidden_states = modulate(self.norm_final(hidden_states), shift, scale)
371
+ return self.linear(hidden_states)
372
+
373
+
374
+ def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False, extra_tokens: int = 0):
375
+ grid_h = np.arange(grid_size, dtype=np.float32)
376
+ grid_w = np.arange(grid_size, dtype=np.float32)
377
+ grid = np.meshgrid(grid_w, grid_h)
378
+ grid = np.stack(grid, axis=0).reshape([2, 1, grid_size, grid_size])
379
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
380
+ if cls_token and extra_tokens > 0:
381
+ pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
382
+ return pos_embed
383
+
384
+
385
+ def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray):
386
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
387
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
388
+ return np.concatenate([emb_h, emb_w], axis=1)
389
+
390
+
391
+ def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray):
392
+ omega = np.arange(embed_dim // 2, dtype=np.float64)
393
+ omega /= embed_dim / 2.0
394
+ omega = 1.0 / 10000**omega
395
+ pos = pos.reshape(-1)
396
+ out = np.einsum("m,d->md", pos, omega)
397
+ return np.concatenate([np.sin(out), np.cos(out)], axis=1)
398
+
399
+
400
+ class DiTMoETransformer2DModel(ModelMixin, ConfigMixin):
401
+ config_name = "config.json"
402
+
403
+ @register_to_config
404
+ def __init__(
405
+ self,
406
+ input_size: int = 32,
407
+ patch_size: int = 2,
408
+ in_channels: int = 4,
409
+ hidden_size: int = 1152,
410
+ depth: int = 28,
411
+ num_heads: int = 16,
412
+ mlp_ratio: float = 4.0,
413
+ class_dropout_prob: float = 0.1,
414
+ num_classes: int = 1000,
415
+ num_experts: int = 8,
416
+ num_experts_per_tok: int = 2,
417
+ pretraining_tp: int = 2,
418
+ learn_sigma: bool = True,
419
+ use_flash_attn: bool = False,
420
+ ):
421
+ super().__init__()
422
+ self.learn_sigma = learn_sigma
423
+ self.in_channels = in_channels
424
+ self.out_channels = in_channels * 2 if learn_sigma else in_channels
425
+ self.patch_size = patch_size
426
+ self.num_heads = num_heads
427
+ self.num_classes = num_classes
428
+
429
+ self.x_embedder = DiTMoEPatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
430
+ self.t_embedder = DiTMoETimestepEmbedder(hidden_size)
431
+ self.y_embedder = DiTMoELabelEmbedder(num_classes, hidden_size, class_dropout_prob)
432
+ self.pos_embed = nn.Parameter(torch.zeros(1, self.x_embedder.num_patches, hidden_size), requires_grad=False)
433
+ self.blocks = nn.ModuleList(
434
+ [
435
+ DiTMoEBlock(
436
+ hidden_size,
437
+ num_heads,
438
+ mlp_ratio=mlp_ratio,
439
+ num_experts=num_experts,
440
+ num_experts_per_tok=num_experts_per_tok,
441
+ pretraining_tp=pretraining_tp,
442
+ use_flash_attn=use_flash_attn,
443
+ )
444
+ for _ in range(depth)
445
+ ]
446
+ )
447
+ self.final_layer = DiTMoEFinalLayer(hidden_size, patch_size, self.out_channels)
448
+ self.initialize_weights()
449
+
450
+ def initialize_weights(self):
451
+ def _basic_init(module):
452
+ if isinstance(module, nn.Linear):
453
+ nn.init.xavier_uniform_(module.weight)
454
+ if module.bias is not None:
455
+ nn.init.constant_(module.bias, 0)
456
+
457
+ self.apply(_basic_init)
458
+ pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches**0.5))
459
+ self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
460
+ w = self.x_embedder.proj.weight.data
461
+ nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
462
+ nn.init.constant_(self.x_embedder.proj.bias, 0)
463
+ nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
464
+ nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
465
+ nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
466
+ for block in self.blocks:
467
+ nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
468
+ nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
469
+ nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
470
+ nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
471
+ nn.init.constant_(self.final_layer.linear.weight, 0)
472
+ nn.init.constant_(self.final_layer.linear.bias, 0)
473
+
474
+ def unpatchify(self, hidden_states: torch.Tensor) -> torch.Tensor:
475
+ channels = self.out_channels
476
+ patch_size = self.patch_size
477
+ height = width = int(hidden_states.shape[1] ** 0.5)
478
+ hidden_states = hidden_states.reshape(hidden_states.shape[0], height, width, patch_size, patch_size, channels)
479
+ hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
480
+ return hidden_states.reshape(hidden_states.shape[0], channels, height * patch_size, width * patch_size)
481
+
482
+ def forward(
483
+ self,
484
+ hidden_states: torch.Tensor,
485
+ timestep: Union[torch.Tensor, float, int],
486
+ class_labels: torch.LongTensor,
487
+ return_dict: bool = True,
488
+ ) -> Union[DiTMoETransformer2DModelOutput, torch.Tensor]:
489
+ if not torch.is_tensor(timestep):
490
+ timestep = torch.tensor([timestep], device=hidden_states.device, dtype=hidden_states.dtype)
491
+ timestep = timestep.to(device=hidden_states.device, dtype=torch.float32)
492
+ if timestep.ndim == 0:
493
+ timestep = timestep.unsqueeze(0)
494
+ if timestep.shape[0] == 1 and hidden_states.shape[0] > 1:
495
+ timestep = timestep.expand(hidden_states.shape[0])
496
+
497
+ class_labels = class_labels.to(device=hidden_states.device, dtype=torch.long).reshape(-1)
498
+ hidden_states = self.x_embedder(hidden_states) + self.pos_embed
499
+ conditioning = self.t_embedder(timestep) + self.y_embedder(class_labels, train=self.training)
500
+ for block in self.blocks:
501
+ hidden_states = block(hidden_states, conditioning)
502
+ hidden_states = self.final_layer(hidden_states, conditioning)
503
+ hidden_states = self.unpatchify(hidden_states)
504
+ if not return_dict:
505
+ return hidden_states
506
+ return DiTMoETransformer2DModelOutput(sample=hidden_states)
507
+
508
+ def split_velocity(self, sample: torch.Tensor) -> torch.Tensor:
509
+ if self.learn_sigma:
510
+ sample, _ = sample.chunk(2, dim=1)
511
+ return sample
DiT-MoE-B-8E2A/vae/config.json ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "AutoencoderKL",
3
+ "_diffusers_version": "0.4.2",
4
+ "act_fn": "silu",
5
+ "block_out_channels": [
6
+ 128,
7
+ 256,
8
+ 512,
9
+ 512
10
+ ],
11
+ "down_block_types": [
12
+ "DownEncoderBlock2D",
13
+ "DownEncoderBlock2D",
14
+ "DownEncoderBlock2D",
15
+ "DownEncoderBlock2D"
16
+ ],
17
+ "in_channels": 3,
18
+ "latent_channels": 4,
19
+ "layers_per_block": 2,
20
+ "norm_num_groups": 32,
21
+ "out_channels": 3,
22
+ "sample_size": 256,
23
+ "up_block_types": [
24
+ "UpDecoderBlock2D",
25
+ "UpDecoderBlock2D",
26
+ "UpDecoderBlock2D",
27
+ "UpDecoderBlock2D"
28
+ ]
29
+ }
DiT-MoE-B-8E2A/vae/diffusion_pytorch_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f3944a3bb4da8058dd3d48327d8da8d875d5ea780833ade1b095bc4049536bac
3
+ size 334643252
DiT-MoE-S-8E2A/README.md CHANGED
@@ -39,6 +39,7 @@ Each variant keeps an English `id2label` map in `model_index.json` (DiT-style).
39
  | Sampler | DDIM |
40
  | Steps | 50 |
41
  | CFG scale | 4.0 |
 
42
  | VAE | `stabilityai/sd-vae-ft-mse` (bundled under `vae/`) |
43
 
44
  ```python
@@ -52,7 +53,7 @@ pipe = DiffusionPipeline.from_pretrained(
52
  local_files_only=True,
53
  custom_pipeline=str(model_dir / "pipeline.py"),
54
  trust_remote_code=True,
55
- torch_dtype=torch.float16,
56
  )
57
  pipe.to("cuda")
58
 
 
39
  | Sampler | DDIM |
40
  | Steps | 50 |
41
  | CFG scale | 4.0 |
42
+ | Dtype | `bfloat16` (recommended on Ampere+) |
43
  | VAE | `stabilityai/sd-vae-ft-mse` (bundled under `vae/`) |
44
 
45
  ```python
 
53
  local_files_only=True,
54
  custom_pipeline=str(model_dir / "pipeline.py"),
55
  trust_remote_code=True,
56
+ torch_dtype=torch.bfloat16,
57
  )
58
  pipe.to("cuda")
59
 
DiT-MoE-S-8E2A/pipeline.py CHANGED
@@ -42,7 +42,7 @@ EXAMPLE_DOC_STRING = """
42
  ... local_files_only=True,
43
  ... custom_pipeline=str(model_dir / "pipeline.py"),
44
  ... trust_remote_code=True,
45
- ... torch_dtype=torch.float16,
46
  ... )
47
  >>> pipe.to("cuda")
48
 
 
42
  ... local_files_only=True,
43
  ... custom_pipeline=str(model_dir / "pipeline.py"),
44
  ... trust_remote_code=True,
45
+ ... torch_dtype=torch.bfloat16,
46
  ... )
47
  >>> pipe.to("cuda")
48
 
README.md CHANGED
@@ -34,7 +34,7 @@ Each variant keeps an English `id2label` map directly in its own `model_index.js
34
  | Checkpoint | Path | Resolution | Sampler |
35
  | --- | --- | --- | --- |
36
  | DiT-MoE-S/2-8E2A | `./DiT-MoE-S-8E2A` | 256×256 | DDIM |
37
- | DiT-MoE-B/2-8E2A | *(convert with same script)* | 256×256 | DDIM |
38
  | DiT-MoE-XL/2-8E2A | *(convert with `--rectified-flow`)* | 256×256 | RF |
39
  | DiT-MoE-G/2-16E2A | *(convert with `--rectified-flow --num-experts 16`)* | 256×256 | RF |
40
 
@@ -56,6 +56,8 @@ python scripts/convert_dit_moe_to_diffusers.py \
56
 
57
  ## Inference
58
 
 
 
59
  ```python
60
  from pathlib import Path
61
  import torch
@@ -67,7 +69,7 @@ pipe = DiffusionPipeline.from_pretrained(
67
  local_files_only=True,
68
  custom_pipeline=str(model_dir / "pipeline.py"),
69
  trust_remote_code=True,
70
- torch_dtype=torch.float16,
71
  )
72
  pipe.to("cuda")
73
 
 
34
  | Checkpoint | Path | Resolution | Sampler |
35
  | --- | --- | --- | --- |
36
  | DiT-MoE-S/2-8E2A | `./DiT-MoE-S-8E2A` | 256×256 | DDIM |
37
+ | DiT-MoE-B/2-8E2A | `./DiT-MoE-B-8E2A` | 256×256 | DDIM |
38
  | DiT-MoE-XL/2-8E2A | *(convert with `--rectified-flow`)* | 256×256 | RF |
39
  | DiT-MoE-G/2-16E2A | *(convert with `--rectified-flow --num-experts 16`)* | 256×256 | RF |
40
 
 
56
 
57
  ## Inference
58
 
59
+ Use `torch.bfloat16` on Ampere+ GPUs (default in examples and `sample_dit_moe.py`).
60
+
61
  ```python
62
  from pathlib import Path
63
  import torch
 
69
  local_files_only=True,
70
  custom_pipeline=str(model_dir / "pipeline.py"),
71
  trust_remote_code=True,
72
+ torch_dtype=torch.bfloat16,
73
  )
74
  pipe.to("cuda")
75