Upload folder using huggingface_hub

.gitattributes

@@ -35,4 +35,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 demo.png filter=lfs diff=lfs merge=lfs -text
 demo_images/jit_h32_test_inference.png filter=lfs diff=lfs merge=lfs -text
-demo_images/jit_h32_final_test.png filter=lfs diff=lfs merge=lfs -text

JiT-B-16/model_index.json

@@ -5,11 +5,1013 @@
   ],
   "_diffusers_version": "0.36.0",
   "scheduler": [
-    "scheduling_jit",
-    "JiTScheduler"
+    "diffusers",
+    "FlowMatchHeunDiscreteScheduler"
   ],
   "transformer": [
     "jit_transformer_2d",
     "JiTTransformer2DModel"
-  ]
+  ],
+  "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"
+  }
 }

JiT-B-16/pipeline.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import importlib
 import json
 import sys
@@ -23,6 +21,7 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchHeunDiscreteScheduler, KarrasDiffusionSchedulers
 from diffusers.utils.torch_utils import randn_tensor
 
 
@@ -39,12 +38,10 @@ class JiTPipeline(DiffusionPipeline):
     Parameters:
         transformer ([`JiTTransformer2DModel`]):
             A class-conditioned `JiTTransformer2DModel` to denoise the images.
-        scheduler ([`JiTScheduler`]):
-            Manual JiT flow-matching scheduler (linear `t in [0, 1]`, Heun or Euler).
+        scheduler ([`KarrasDiffusionSchedulers`] or [`FlowMatchHeunDiscreteScheduler`]):
+            Diffusers scheduler interface for JiT generation (defaults to `FlowMatchHeunDiscreteScheduler(shift=4.0)`).
         id2label (`dict[int, str]`, *optional*):
             ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
-        id2label_cn (`dict[int, str]`, *optional*):
-            ImageNet class id to Chinese label mapping. Values may contain comma-separated synonyms.
     """
 
     model_cpu_offload_seq = "transformer"
@@ -71,7 +68,7 @@ class JiTPipeline(DiffusionPipeline):
 
             hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
             if subfolder:
-                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**", "labels/**"])
+                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
             cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
             variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
         else:
@@ -82,6 +79,7 @@ class JiTPipeline(DiffusionPipeline):
             if subfolder:
                 variant = variant / subfolder
 
+        id2label_override = kwargs.pop("id2label", None)
         model_kwargs = dict(kwargs)
         inserted: List[str] = []
 
@@ -103,19 +101,22 @@ class JiTPipeline(DiffusionPipeline):
 
         try:
             transformer = _load_component("transformer", "jit_transformer_2d", "JiTTransformer2DModel")
-            scheduler = _load_component("scheduler", "scheduling_jit", "JiTScheduler")
+            try:
+                scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained(str(variant), subfolder="scheduler")
+            except Exception:
+                scheduler = FlowMatchHeunDiscreteScheduler(shift=4.0)
 
             if transformer is None:
                 raise ValueError(f"No loadable transformer found under {variant}")
 
             variant_path = str(variant)
-            id2label, id2label_cn = cls._load_labels_for_variant(variant_path)
+            model_index_path = variant / "model_index.json"
+            id2label = id2label_override or cls._read_id2label_from_model_index(model_index_path)
 
             pipe = cls(
                 transformer=transformer,
                 scheduler=scheduler,
                 id2label=id2label,
-                id2label_cn=id2label_cn,
             )
             if variant_path and hasattr(pipe, "register_to_config"):
                 pipe.register_to_config(_name_or_path=variant_path)
@@ -128,58 +129,31 @@ class JiTPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler,
-        id2label: Optional[Dict[int, str]] = None,
-        id2label_cn: Optional[Dict[int, str]] = None,
+        scheduler: FlowMatchHeunDiscreteScheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
     ):
         super().__init__()
+        scheduler = scheduler or FlowMatchHeunDiscreteScheduler(shift=4.0)
         self.register_modules(transformer=transformer, scheduler=scheduler)
 
-        self._id2label = id2label or {}
-        self._id2label_cn = id2label_cn or {}
+        self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
-        self.labels_cn = self._build_label2id(self._id2label_cn)
-
-    def _ensure_labels_loaded(self) -> None:
-        if self._id2label or self._id2label_cn:
-            return
-        loaded_en, loaded_cn = self._load_labels_for_variant(getattr(self.config, "_name_or_path", None))
-        if loaded_en:
-            self._id2label = loaded_en
-            self.labels = self._build_label2id(self._id2label)
-        if loaded_cn:
-            self._id2label_cn = loaded_cn
-            self.labels_cn = self._build_label2id(self._id2label_cn)
 
     @staticmethod
-    def _labels_dir_for_variant(variant_path: Optional[str]) -> Optional[Path]:
-        if not variant_path:
-            return None
-        variant_dir = Path(variant_path).resolve()
-        labels_dir = variant_dir.parent / "labels"
-        return labels_dir if labels_dir.is_dir() else None
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
 
     @staticmethod
-    def _read_id2label(labels_dir: Path, lang: str = "en") -> Dict[int, str]:
-        filename = "id2label_en.json" if lang == "en" else "id2label_cn.json"
-        path = labels_dir / filename
-        if not path.exists():
-            raise FileNotFoundError(path)
-        raw = json.loads(path.read_text(encoding="utf-8"))
-        return {int(key): value for key, value in raw.items()}
-
-    @classmethod
-    def _load_labels_for_variant(
-        cls,
-        variant_path: Optional[str],
-    ) -> Tuple[Optional[Dict[int, str]], Optional[Dict[int, str]]]:
-        labels_dir = cls._labels_dir_for_variant(variant_path)
-        if labels_dir is None:
-            return None, None
-        try:
-            return cls._read_id2label(labels_dir, "en"), cls._read_id2label(labels_dir, "cn")
-        except FileNotFoundError:
-            return None, None
+    def _read_id2label_from_model_index(model_index_path: Path) -> Dict[int, str]:
+        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]:
@@ -194,35 +168,19 @@ class JiTPipeline(DiffusionPipeline):
     @property
     def id2label(self) -> Dict[int, str]:
         """ImageNet class id to English label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
         return self._id2label
 
-    @property
-    def id2label_cn(self) -> Dict[int, str]:
-        """ImageNet class id to Chinese label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
-        return self._id2label_cn
-
-    def get_label_ids(self, label: Union[str, List[str]], lang: str = "en") -> List[int]:
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
         r"""
         Map ImageNet label strings to class ids.
 
         Args:
             label (`str` or `list[str]`):
-                One or more label strings. Each string must match a synonym in `id2label` (English)
-                or `id2label_cn` (Chinese).
-            lang (`str`, *optional*, defaults to `"en"`):
-                `"en"` uses English synonyms; `"cn"` uses Chinese synonyms.
+                One or more English label strings. Each string must match a synonym in `id2label`.
         """
-        if lang not in ("en", "cn"):
-            raise ValueError(f"`lang` must be 'en' or 'cn', got {lang!r}.")
-
-        self._ensure_labels_loaded()
-        label2id = self.labels if lang == "en" else self.labels_cn
+        label2id = self.labels
         if not label2id:
-            raise ValueError(
-                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
-            )
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
 
         if isinstance(label, str):
             label = [label]
@@ -231,7 +189,7 @@ class JiTPipeline(DiffusionPipeline):
         if missing:
             preview = ", ".join(list(label2id.keys())[:8])
             raise ValueError(
-                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+                f"Unknown English label(s): {missing}. Example valid labels: {preview}, ..."
             )
         return [label2id[item] for item in label]
 
@@ -246,115 +204,10 @@ class JiTPipeline(DiffusionPipeline):
             return self.get_label_ids(class_labels)
 
         if class_labels and isinstance(class_labels[0], str):
-            self._ensure_labels_loaded()
-            if all(label in self.labels for label in class_labels):
-                return self.get_label_ids(class_labels, lang="en")
-            if all(label in self.labels_cn for label in class_labels):
-                return self.get_label_ids(class_labels, lang="cn")
-            raise ValueError(
-                "Could not resolve string `class_labels`. Use English synonyms from `pipe.labels` "
-                "or Chinese synonyms from `pipe.labels_cn`."
-            )
+            return self.get_label_ids(class_labels)
 
         return list(class_labels)
 
-    def _predict_velocity(
-        self,
-        z_value: torch.Tensor,
-        t: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-    ) -> torch.Tensor:
-        t = torch.as_tensor(t, device=z_value.device, dtype=z_value.dtype)
-        if do_classifier_free_guidance:
-            z_in = torch.cat([z_value, z_value], dim=0)
-            labels = torch.cat([class_labels, class_null], dim=0)
-        else:
-            z_in = z_value
-            labels = class_labels
-
-        t_batch = t.flatten().expand(z_in.shape[0])
-        x_pred = self.transformer(z_in, timestep=t_batch, class_labels=labels).sample
-        v = self.scheduler.velocity_from_prediction(z_in, x_pred, t)
-
-        if not do_classifier_free_guidance:
-            return v
-
-        v_cond, v_uncond = v.chunk(2, dim=0)
-        interval_mask = t < guidance_interval_max
-        if guidance_interval_min != 0.0:
-            interval_mask = interval_mask & (t > guidance_interval_min)
-        scale = torch.where(
-            interval_mask,
-            torch.tensor(guidance_scale, device=z_value.device, dtype=z_value.dtype),
-            torch.tensor(1.0, device=z_value.device, dtype=z_value.dtype),
-        )
-        return v_uncond + scale * (v_cond - v_uncond)
-
-    def _run_sampler(
-        self,
-        latents: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        num_inference_steps: int,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-        sampling_method: str,
-    ) -> torch.Tensor:
-        device = latents.device
-        self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampling_method)
-        timesteps = self.scheduler.timesteps
-
-        for i in self.progress_bar(range(num_inference_steps - 1)):
-            t = timesteps[i]
-            t_next = timesteps[i + 1]
-            v = self._predict_velocity(
-                latents,
-                t,
-                class_labels,
-                class_null,
-                do_classifier_free_guidance,
-                guidance_scale,
-                guidance_interval_min,
-                guidance_interval_max,
-            )
-
-            if sampling_method == "heun":
-                latents_euler = latents + (t_next - t) * v
-                v_next = self._predict_velocity(
-                    latents_euler,
-                    t_next,
-                    class_labels,
-                    class_null,
-                    do_classifier_free_guidance,
-                    guidance_scale,
-                    guidance_interval_min,
-                    guidance_interval_max,
-                )
-                latents = self.scheduler.step(v, t, latents, model_output_next=v_next).prev_sample
-            else:
-                latents = self.scheduler.step(v, t, latents).prev_sample
-
-        t = timesteps[-2]
-        t_next = timesteps[-1]
-        v = self._predict_velocity(
-            latents,
-            t,
-            class_labels,
-            class_null,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-        )
-        return latents + (t_next - t) * v
-
     @torch.inference_mode()
     def __call__(
         self,
@@ -363,10 +216,12 @@ class JiTPipeline(DiffusionPipeline):
         guidance_interval_min: float = 0.1,
         guidance_interval_max: float = 1.0,
         noise_scale: Optional[float] = None,
-        t_eps: Optional[float] = None,
-        sampling_method: Optional[str] = None,
+        t_eps: float = 5e-2,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         num_inference_steps: int = 50,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        interpolate_pos_encoding: bool = True,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
     ) -> Union[ImagePipelineOutput, Tuple]:
@@ -375,7 +230,7 @@ class JiTPipeline(DiffusionPipeline):
 
         Args:
             class_labels (`int`, `str`, `list[int]`, or `list[str]`):
-                ImageNet class indices or human-readable label strings (English or Chinese).
+                ImageNet class indices or human-readable English label strings.
             guidance_scale (`float`, *optional*):
                 Classifier-free guidance scale. CFG is active when `guidance_scale > 1.0`.
             guidance_interval_min (`float`, defaults to `0.1`):
@@ -384,10 +239,8 @@ class JiTPipeline(DiffusionPipeline):
                 Upper bound of the CFG interval in flow time.
             noise_scale (`float`, *optional*):
                 Initial Gaussian noise scale (`1.0` for 256px, `2.0` for 512px by default).
-            t_eps (`float`, *optional*):
-                Epsilon clamp for the `1 - t` denominator (scheduler config by default).
-            sampling_method (`str`, *optional*):
-                `"heun"` or `"euler"`. Defaults to the scheduler config (`heun`).
+            t_eps (`float`, defaults to `5e-2`):
+                Epsilon clamp for the `1 - t` denominator, matching JiT source defaults.
             generator (`torch.Generator`, *optional*):
                 RNG for reproducibility.
             num_inference_steps (`int`, defaults to `50`):
@@ -397,31 +250,34 @@ class JiTPipeline(DiffusionPipeline):
             return_dict (`bool`, *optional*, defaults to `True`):
                 Return [`ImagePipelineOutput`] if True.
         """
-        solver = sampling_method or self.scheduler.config.solver
-        if solver not in {"heun", "euler"}:
-            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
         if num_inference_steps < 2:
             raise ValueError("num_inference_steps must be >= 2.")
 
-        if t_eps is not None:
-            self.scheduler.register_to_config(t_eps=t_eps)
-
         class_label_ids = self._normalize_class_labels(class_labels)
         do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
 
         batch_size = len(class_label_ids)
         image_size = int(self.transformer.config.sample_size)
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
         channels = int(self.transformer.config.in_channels)
         null_class_val = int(self.transformer.config.num_classes)
 
         if guidance_scale is None:
             guidance_scale = 1.0
         if noise_scale is None:
-            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(image_size, 1.0)
+            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)
 
         latents = (
             randn_tensor(
-                shape=(batch_size, channels, image_size, image_size),
+                shape=(batch_size, channels, height, width),
                 generator=generator,
                 device=self._execution_device,
                 dtype=self.transformer.dtype,
@@ -433,17 +289,47 @@ class JiTPipeline(DiffusionPipeline):
         class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
         class_null = torch.full_like(class_labels_t, null_class_val)
 
-        latents = self._run_sampler(
-            latents,
-            class_labels_t,
-            class_null,
-            num_inference_steps,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-            solver,
-        )
+        if do_classifier_free_guidance:
+            class_labels_input = torch.cat([class_labels_t, class_null], dim=0)
+        else:
+            class_labels_input = class_labels_t
+
+        self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
+        for t in self.progress_bar(self.scheduler.timesteps):
+            step_index = self.scheduler.index_for_timestep(t, self.scheduler.timesteps)
+            sigma = self.scheduler.sigmas[step_index].to(device=latents.device, dtype=latents.dtype)
+            sigma = sigma.clamp_min(t_eps)
+            t_flow = (1.0 - sigma).clamp(0.0, 1.0)
+
+            if do_classifier_free_guidance:
+                latent_model_input = torch.cat([latents, latents], dim=0)
+            else:
+                latent_model_input = latents
+
+            timesteps = t_flow.flatten().expand(latent_model_input.shape[0])
+            x_pred = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=class_labels_input,
+                interpolate_pos_encoding=interpolate_pos_encoding,
+            ).sample
+
+            if do_classifier_free_guidance:
+                x_cond, x_uncond = x_pred.chunk(2, dim=0)
+                interval_mask = t_flow < guidance_interval_max
+                if guidance_interval_min != 0.0:
+                    interval_mask = interval_mask & (t_flow > guidance_interval_min)
+                scale = torch.where(
+                    interval_mask,
+                    torch.tensor(guidance_scale, device=latents.device, dtype=latents.dtype),
+                    torch.tensor(1.0, device=latents.device, dtype=latents.dtype),
+                )
+                x_pred = x_uncond + scale * (x_cond - x_uncond)
+
+            sigma = sigma.reshape(*([1] * (latents.ndim - 1)))
+            # JiT predicts x0; scheduler integrates in sigma space: dz/dsigma = -(x0 - z) / sigma.
+            model_output = -(x_pred - latents) / sigma
+            latents = self.scheduler.step(model_output, t, latents).prev_sample
 
         images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
         if output_type == "pt":

JiT-B-16/scheduler/scheduler_config.json

@@ -1,7 +1,6 @@
 {
-  "_class_name": "JiTScheduler",
+  "_class_name": "FlowMatchHeunDiscreteScheduler",
   "_diffusers_version": "0.36.0",
   "num_train_timesteps": 1000,
-  "t_eps": 0.05,
-  "solver": "heun"
+  "shift": 4.0
 }

JiT-B-16/transformer/jit_transformer_2d.py

@@ -68,38 +68,58 @@ class JiTRotaryEmbedding(nn.Module):
         num_cls_token=0,
     ):
         super().__init__()
-        if custom_freqs is not None:
-            freqs = custom_freqs
-        else:
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-
+        self.dim = dim
+        self.pt_seq_len = pt_seq_len
+        self.theta = theta
+        self.num_cls_token = num_cls_token
+        self.custom_freqs = custom_freqs
         if ft_seq_len is None:
             ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = freqs.repeat_interleave(2, dim=-1)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])  # [N_img, D]
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-
-            # prepend in-context cls token
-            _, D = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, D, dtype=cos_img.dtype)
-            sin_pad = torch.zeros(num_cls_token, D, dtype=sin_img.dtype)
-
-            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
-            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
+        self._cached_hw = None
+        cos, sin = self._build_freqs(ft_seq_len, ft_seq_len, device=torch.device("cpu"))
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+        self._cached_hw = (ft_seq_len, ft_seq_len)
+
+    def _build_freqs(self, height, width, device):
+        if self.custom_freqs is not None:
+            freqs = self.custom_freqs.to(device=device, dtype=torch.float32)
         else:
-            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
-            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
-
-    def forward(self, t):
+            freqs = 1.0 / (
+                self.theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)[: (self.dim // 2)] / self.dim)
+            )
+
+        t_h = torch.arange(height, device=device, dtype=torch.float32) / height * self.pt_seq_len
+        t_w = torch.arange(width, device=device, dtype=torch.float32) / width * self.pt_seq_len
+        freqs_h = torch.einsum("..., f -> ... f", t_h, freqs).repeat_interleave(2, dim=-1)
+        freqs_w = torch.einsum("..., f -> ... f", t_w, freqs).repeat_interleave(2, dim=-1)
+        freqs_2d = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
+        freqs_flat = freqs_2d.view(-1, freqs_2d.shape[-1])
+        cos_img = freqs_flat.cos()
+        sin_img = freqs_flat.sin()
+        if self.num_cls_token > 0:
+            _, dim_freq = cos_img.shape
+            cos_pad = torch.ones(self.num_cls_token, dim_freq, dtype=cos_img.dtype, device=device)
+            sin_pad = torch.zeros(self.num_cls_token, dim_freq, dtype=sin_img.dtype, device=device)
+            cos_img = torch.cat([cos_pad, cos_img], dim=0)
+            sin_img = torch.cat([sin_pad, sin_img], dim=0)
+        return cos_img, sin_img
+
+    def forward(self, t, height=None, width=None):
         # Applied on (batch, seq_len, heads, head_dim) tensors from attention.
         seq_len = t.shape[1]
+        if height is None or width is None:
+            image_tokens = seq_len - self.num_cls_token
+            size = int(image_tokens**0.5)
+            if size * size != image_tokens:
+                raise ValueError(
+                    f"Cannot infer square token grid from sequence length {seq_len} with {self.num_cls_token} class tokens."
+                )
+            height = size
+            width = size
+        if self._cached_hw != (height, width) or self.freqs_cos.device != t.device:
+            self.freqs_cos, self.freqs_sin = self._build_freqs(height, width, device=t.device)
+            self._cached_hw = (height, width)
         freqs_cos = self.freqs_cos[:seq_len].to(t.dtype)
         freqs_sin = self.freqs_sin[:seq_len].to(t.dtype)
 
@@ -195,7 +215,7 @@ class JiTAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
 
-    def forward(self, x, rope=None):
+    def forward(self, x, rope=None, grid_height=None, grid_width=None):
         B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
@@ -206,8 +226,8 @@ class JiTAttention(nn.Module):
         if rope is not None:
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
-            q = rope(q)
-            k = rope(k)
+            q = rope(q, height=grid_height, width=grid_width)
+            k = rope(k, height=grid_height, width=grid_width)
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
 
@@ -254,7 +274,7 @@ class JiTBlock(nn.Module):
         self.act = nn.SiLU()
         self.adaLN_modulation = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
 
-    def forward(self, x, c, feat_rope=None):
+    def forward(self, x, c, feat_rope=None, grid_height=None, grid_width=None):
         # Apply activation
         c = self.act(c)
 
@@ -263,7 +283,7 @@ class JiTBlock(nn.Module):
         # Attention block
         norm_x = self.norm1(x)
         modulated_x = modulate(norm_x, shift_msa, scale_msa)
-        attn_out = self.attn(modulated_x, rope=feat_rope)
+        attn_out = self.attn(modulated_x, rope=feat_rope, grid_height=grid_height, grid_width=grid_width)
         x = x + gate_msa.unsqueeze(1) * attn_out
 
         # MLP block
@@ -437,11 +457,30 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         self.act_final = nn.SiLU()
         self.adaLN_modulation_final = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
 
+    def _get_patch_grid(self, hidden_states):
+        height, width = hidden_states.shape[-2:]
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"Input size {(height, width)} must be divisible by patch_size={self.patch_size}."
+            )
+        return height // self.patch_size, width // self.patch_size
+
+    def _interpolate_pos_encoding(self, tokens, grid_height, grid_width):
+        num_tokens = grid_height * grid_width
+        if self.pos_embed.shape[1] == num_tokens:
+            return self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+        base_size = int(self.pos_embed.shape[1] ** 0.5)
+        pos_embed = self.pos_embed.reshape(1, base_size, base_size, self.hidden_size).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(grid_height, grid_width), mode="bicubic", align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(1, num_tokens, self.hidden_size)
+        return pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         class_labels: torch.LongTensor,
+        interpolate_pos_encoding: bool = True,
         return_dict: bool = True,
     ):
 
@@ -454,8 +493,19 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         c = t_emb + y_emb
 
         # Patch Embed
+        grid_height, grid_width = self._get_patch_grid(hidden_states)
         x = self.x_embedder(hidden_states)
-        x = x + self.pos_embed.to(x.dtype)
+        if interpolate_pos_encoding:
+            pos_embed = self._interpolate_pos_encoding(x, grid_height, grid_width)
+        else:
+            expected_tokens = grid_height * grid_width
+            if self.pos_embed.shape[1] != expected_tokens:
+                raise ValueError(
+                    f"pos_embed token count {self.pos_embed.shape[1]} does not match input token count {expected_tokens}. "
+                    "Enable interpolate_pos_encoding for dynamic resolutions."
+                )
+            pos_embed = self.pos_embed.to(device=x.device, dtype=x.dtype)
+        x = x + pos_embed
 
         # Blocks
         for i, block in enumerate(self.blocks):
@@ -467,15 +517,23 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
             rope = self.feat_rope if i < self.in_context_start else self.feat_rope_incontext
 
             if self.training and self.gradient_checkpointing:
+                def custom_forward(current_x, current_c):
+                    return block(
+                        current_x,
+                        current_c,
+                        feat_rope=rope,
+                        grid_height=grid_height,
+                        grid_width=grid_width,
+                    )
+
                 x = torch.utils.checkpoint.checkpoint(
-                    block,
+                    custom_forward,
                     x,
                     c,
-                    rope,
                     use_reentrant=False,
                 )
             else:
-                x = block(x, c, feat_rope=rope)
+                x = block(x, c, feat_rope=rope, grid_height=grid_height, grid_width=grid_width)
 
         # Slice off in-context tokens
         if self.in_context_len > 0:
@@ -489,10 +547,11 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         x = self.linear_final(x)
 
         # Unpatchify
-        h = w = int(x.shape[1] ** 0.5)
-        x = x.reshape(shape=(x.shape[0], h, w, self.patch_size, self.patch_size, self.out_channels))
+        x = x.reshape(shape=(x.shape[0], grid_height, grid_width, self.patch_size, self.patch_size, self.out_channels))
         x = torch.einsum("nhwpqc->nchpwq", x)
-        output = x.reshape(shape=(x.shape[0], self.out_channels, h * self.patch_size, w * self.patch_size))
+        output = x.reshape(
+            shape=(x.shape[0], self.out_channels, grid_height * self.patch_size, grid_width * self.patch_size)
+        )
 
         if not return_dict:
             return (output,)

JiT-B-32/model_index.json

@@ -5,11 +5,1013 @@
   ],
   "_diffusers_version": "0.36.0",
   "scheduler": [
-    "scheduling_jit",
-    "JiTScheduler"
+    "diffusers",
+    "FlowMatchHeunDiscreteScheduler"
   ],
   "transformer": [
     "jit_transformer_2d",
     "JiTTransformer2DModel"
-  ]
+  ],
+  "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"
+  }
 }

JiT-B-32/pipeline.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import importlib
 import json
 import sys
@@ -23,6 +21,7 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchHeunDiscreteScheduler, KarrasDiffusionSchedulers
 from diffusers.utils.torch_utils import randn_tensor
 
 
@@ -39,12 +38,10 @@ class JiTPipeline(DiffusionPipeline):
     Parameters:
         transformer ([`JiTTransformer2DModel`]):
             A class-conditioned `JiTTransformer2DModel` to denoise the images.
-        scheduler ([`JiTScheduler`]):
-            Manual JiT flow-matching scheduler (linear `t in [0, 1]`, Heun or Euler).
+        scheduler ([`KarrasDiffusionSchedulers`] or [`FlowMatchHeunDiscreteScheduler`]):
+            Diffusers scheduler interface for JiT generation (defaults to `FlowMatchHeunDiscreteScheduler(shift=4.0)`).
         id2label (`dict[int, str]`, *optional*):
             ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
-        id2label_cn (`dict[int, str]`, *optional*):
-            ImageNet class id to Chinese label mapping. Values may contain comma-separated synonyms.
     """
 
     model_cpu_offload_seq = "transformer"
@@ -71,7 +68,7 @@ class JiTPipeline(DiffusionPipeline):
 
             hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
             if subfolder:
-                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**", "labels/**"])
+                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
             cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
             variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
         else:
@@ -82,6 +79,7 @@ class JiTPipeline(DiffusionPipeline):
             if subfolder:
                 variant = variant / subfolder
 
+        id2label_override = kwargs.pop("id2label", None)
         model_kwargs = dict(kwargs)
         inserted: List[str] = []
 
@@ -103,19 +101,22 @@ class JiTPipeline(DiffusionPipeline):
 
         try:
             transformer = _load_component("transformer", "jit_transformer_2d", "JiTTransformer2DModel")
-            scheduler = _load_component("scheduler", "scheduling_jit", "JiTScheduler")
+            try:
+                scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained(str(variant), subfolder="scheduler")
+            except Exception:
+                scheduler = FlowMatchHeunDiscreteScheduler(shift=4.0)
 
             if transformer is None:
                 raise ValueError(f"No loadable transformer found under {variant}")
 
             variant_path = str(variant)
-            id2label, id2label_cn = cls._load_labels_for_variant(variant_path)
+            model_index_path = variant / "model_index.json"
+            id2label = id2label_override or cls._read_id2label_from_model_index(model_index_path)
 
             pipe = cls(
                 transformer=transformer,
                 scheduler=scheduler,
                 id2label=id2label,
-                id2label_cn=id2label_cn,
             )
             if variant_path and hasattr(pipe, "register_to_config"):
                 pipe.register_to_config(_name_or_path=variant_path)
@@ -128,58 +129,31 @@ class JiTPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler,
-        id2label: Optional[Dict[int, str]] = None,
-        id2label_cn: Optional[Dict[int, str]] = None,
+        scheduler: FlowMatchHeunDiscreteScheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
     ):
         super().__init__()
+        scheduler = scheduler or FlowMatchHeunDiscreteScheduler(shift=4.0)
         self.register_modules(transformer=transformer, scheduler=scheduler)
 
-        self._id2label = id2label or {}
-        self._id2label_cn = id2label_cn or {}
+        self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
-        self.labels_cn = self._build_label2id(self._id2label_cn)
-
-    def _ensure_labels_loaded(self) -> None:
-        if self._id2label or self._id2label_cn:
-            return
-        loaded_en, loaded_cn = self._load_labels_for_variant(getattr(self.config, "_name_or_path", None))
-        if loaded_en:
-            self._id2label = loaded_en
-            self.labels = self._build_label2id(self._id2label)
-        if loaded_cn:
-            self._id2label_cn = loaded_cn
-            self.labels_cn = self._build_label2id(self._id2label_cn)
 
     @staticmethod
-    def _labels_dir_for_variant(variant_path: Optional[str]) -> Optional[Path]:
-        if not variant_path:
-            return None
-        variant_dir = Path(variant_path).resolve()
-        labels_dir = variant_dir.parent / "labels"
-        return labels_dir if labels_dir.is_dir() else None
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
 
     @staticmethod
-    def _read_id2label(labels_dir: Path, lang: str = "en") -> Dict[int, str]:
-        filename = "id2label_en.json" if lang == "en" else "id2label_cn.json"
-        path = labels_dir / filename
-        if not path.exists():
-            raise FileNotFoundError(path)
-        raw = json.loads(path.read_text(encoding="utf-8"))
-        return {int(key): value for key, value in raw.items()}
-
-    @classmethod
-    def _load_labels_for_variant(
-        cls,
-        variant_path: Optional[str],
-    ) -> Tuple[Optional[Dict[int, str]], Optional[Dict[int, str]]]:
-        labels_dir = cls._labels_dir_for_variant(variant_path)
-        if labels_dir is None:
-            return None, None
-        try:
-            return cls._read_id2label(labels_dir, "en"), cls._read_id2label(labels_dir, "cn")
-        except FileNotFoundError:
-            return None, None
+    def _read_id2label_from_model_index(model_index_path: Path) -> Dict[int, str]:
+        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]:
@@ -194,35 +168,19 @@ class JiTPipeline(DiffusionPipeline):
     @property
     def id2label(self) -> Dict[int, str]:
         """ImageNet class id to English label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
         return self._id2label
 
-    @property
-    def id2label_cn(self) -> Dict[int, str]:
-        """ImageNet class id to Chinese label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
-        return self._id2label_cn
-
-    def get_label_ids(self, label: Union[str, List[str]], lang: str = "en") -> List[int]:
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
         r"""
         Map ImageNet label strings to class ids.
 
         Args:
             label (`str` or `list[str]`):
-                One or more label strings. Each string must match a synonym in `id2label` (English)
-                or `id2label_cn` (Chinese).
-            lang (`str`, *optional*, defaults to `"en"`):
-                `"en"` uses English synonyms; `"cn"` uses Chinese synonyms.
+                One or more English label strings. Each string must match a synonym in `id2label`.
         """
-        if lang not in ("en", "cn"):
-            raise ValueError(f"`lang` must be 'en' or 'cn', got {lang!r}.")
-
-        self._ensure_labels_loaded()
-        label2id = self.labels if lang == "en" else self.labels_cn
+        label2id = self.labels
         if not label2id:
-            raise ValueError(
-                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
-            )
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
 
         if isinstance(label, str):
             label = [label]
@@ -231,7 +189,7 @@ class JiTPipeline(DiffusionPipeline):
         if missing:
             preview = ", ".join(list(label2id.keys())[:8])
             raise ValueError(
-                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+                f"Unknown English label(s): {missing}. Example valid labels: {preview}, ..."
             )
         return [label2id[item] for item in label]
 
@@ -246,115 +204,10 @@ class JiTPipeline(DiffusionPipeline):
             return self.get_label_ids(class_labels)
 
         if class_labels and isinstance(class_labels[0], str):
-            self._ensure_labels_loaded()
-            if all(label in self.labels for label in class_labels):
-                return self.get_label_ids(class_labels, lang="en")
-            if all(label in self.labels_cn for label in class_labels):
-                return self.get_label_ids(class_labels, lang="cn")
-            raise ValueError(
-                "Could not resolve string `class_labels`. Use English synonyms from `pipe.labels` "
-                "or Chinese synonyms from `pipe.labels_cn`."
-            )
+            return self.get_label_ids(class_labels)
 
         return list(class_labels)
 
-    def _predict_velocity(
-        self,
-        z_value: torch.Tensor,
-        t: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-    ) -> torch.Tensor:
-        t = torch.as_tensor(t, device=z_value.device, dtype=z_value.dtype)
-        if do_classifier_free_guidance:
-            z_in = torch.cat([z_value, z_value], dim=0)
-            labels = torch.cat([class_labels, class_null], dim=0)
-        else:
-            z_in = z_value
-            labels = class_labels
-
-        t_batch = t.flatten().expand(z_in.shape[0])
-        x_pred = self.transformer(z_in, timestep=t_batch, class_labels=labels).sample
-        v = self.scheduler.velocity_from_prediction(z_in, x_pred, t)
-
-        if not do_classifier_free_guidance:
-            return v
-
-        v_cond, v_uncond = v.chunk(2, dim=0)
-        interval_mask = t < guidance_interval_max
-        if guidance_interval_min != 0.0:
-            interval_mask = interval_mask & (t > guidance_interval_min)
-        scale = torch.where(
-            interval_mask,
-            torch.tensor(guidance_scale, device=z_value.device, dtype=z_value.dtype),
-            torch.tensor(1.0, device=z_value.device, dtype=z_value.dtype),
-        )
-        return v_uncond + scale * (v_cond - v_uncond)
-
-    def _run_sampler(
-        self,
-        latents: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        num_inference_steps: int,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-        sampling_method: str,
-    ) -> torch.Tensor:
-        device = latents.device
-        self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampling_method)
-        timesteps = self.scheduler.timesteps
-
-        for i in self.progress_bar(range(num_inference_steps - 1)):
-            t = timesteps[i]
-            t_next = timesteps[i + 1]
-            v = self._predict_velocity(
-                latents,
-                t,
-                class_labels,
-                class_null,
-                do_classifier_free_guidance,
-                guidance_scale,
-                guidance_interval_min,
-                guidance_interval_max,
-            )
-
-            if sampling_method == "heun":
-                latents_euler = latents + (t_next - t) * v
-                v_next = self._predict_velocity(
-                    latents_euler,
-                    t_next,
-                    class_labels,
-                    class_null,
-                    do_classifier_free_guidance,
-                    guidance_scale,
-                    guidance_interval_min,
-                    guidance_interval_max,
-                )
-                latents = self.scheduler.step(v, t, latents, model_output_next=v_next).prev_sample
-            else:
-                latents = self.scheduler.step(v, t, latents).prev_sample
-
-        t = timesteps[-2]
-        t_next = timesteps[-1]
-        v = self._predict_velocity(
-            latents,
-            t,
-            class_labels,
-            class_null,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-        )
-        return latents + (t_next - t) * v
-
     @torch.inference_mode()
     def __call__(
         self,
@@ -363,10 +216,12 @@ class JiTPipeline(DiffusionPipeline):
         guidance_interval_min: float = 0.1,
         guidance_interval_max: float = 1.0,
         noise_scale: Optional[float] = None,
-        t_eps: Optional[float] = None,
-        sampling_method: Optional[str] = None,
+        t_eps: float = 5e-2,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         num_inference_steps: int = 50,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        interpolate_pos_encoding: bool = True,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
     ) -> Union[ImagePipelineOutput, Tuple]:
@@ -375,7 +230,7 @@ class JiTPipeline(DiffusionPipeline):
 
         Args:
             class_labels (`int`, `str`, `list[int]`, or `list[str]`):
-                ImageNet class indices or human-readable label strings (English or Chinese).
+                ImageNet class indices or human-readable English label strings.
             guidance_scale (`float`, *optional*):
                 Classifier-free guidance scale. CFG is active when `guidance_scale > 1.0`.
             guidance_interval_min (`float`, defaults to `0.1`):
@@ -384,10 +239,8 @@ class JiTPipeline(DiffusionPipeline):
                 Upper bound of the CFG interval in flow time.
             noise_scale (`float`, *optional*):
                 Initial Gaussian noise scale (`1.0` for 256px, `2.0` for 512px by default).
-            t_eps (`float`, *optional*):
-                Epsilon clamp for the `1 - t` denominator (scheduler config by default).
-            sampling_method (`str`, *optional*):
-                `"heun"` or `"euler"`. Defaults to the scheduler config (`heun`).
+            t_eps (`float`, defaults to `5e-2`):
+                Epsilon clamp for the `1 - t` denominator, matching JiT source defaults.
             generator (`torch.Generator`, *optional*):
                 RNG for reproducibility.
             num_inference_steps (`int`, defaults to `50`):
@@ -397,31 +250,34 @@ class JiTPipeline(DiffusionPipeline):
             return_dict (`bool`, *optional*, defaults to `True`):
                 Return [`ImagePipelineOutput`] if True.
         """
-        solver = sampling_method or self.scheduler.config.solver
-        if solver not in {"heun", "euler"}:
-            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
         if num_inference_steps < 2:
             raise ValueError("num_inference_steps must be >= 2.")
 
-        if t_eps is not None:
-            self.scheduler.register_to_config(t_eps=t_eps)
-
         class_label_ids = self._normalize_class_labels(class_labels)
         do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
 
         batch_size = len(class_label_ids)
         image_size = int(self.transformer.config.sample_size)
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
         channels = int(self.transformer.config.in_channels)
         null_class_val = int(self.transformer.config.num_classes)
 
         if guidance_scale is None:
             guidance_scale = 1.0
         if noise_scale is None:
-            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(image_size, 1.0)
+            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)
 
         latents = (
             randn_tensor(
-                shape=(batch_size, channels, image_size, image_size),
+                shape=(batch_size, channels, height, width),
                 generator=generator,
                 device=self._execution_device,
                 dtype=self.transformer.dtype,
@@ -433,17 +289,47 @@ class JiTPipeline(DiffusionPipeline):
         class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
         class_null = torch.full_like(class_labels_t, null_class_val)
 
-        latents = self._run_sampler(
-            latents,
-            class_labels_t,
-            class_null,
-            num_inference_steps,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-            solver,
-        )
+        if do_classifier_free_guidance:
+            class_labels_input = torch.cat([class_labels_t, class_null], dim=0)
+        else:
+            class_labels_input = class_labels_t
+
+        self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
+        for t in self.progress_bar(self.scheduler.timesteps):
+            step_index = self.scheduler.index_for_timestep(t, self.scheduler.timesteps)
+            sigma = self.scheduler.sigmas[step_index].to(device=latents.device, dtype=latents.dtype)
+            sigma = sigma.clamp_min(t_eps)
+            t_flow = (1.0 - sigma).clamp(0.0, 1.0)
+
+            if do_classifier_free_guidance:
+                latent_model_input = torch.cat([latents, latents], dim=0)
+            else:
+                latent_model_input = latents
+
+            timesteps = t_flow.flatten().expand(latent_model_input.shape[0])
+            x_pred = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=class_labels_input,
+                interpolate_pos_encoding=interpolate_pos_encoding,
+            ).sample
+
+            if do_classifier_free_guidance:
+                x_cond, x_uncond = x_pred.chunk(2, dim=0)
+                interval_mask = t_flow < guidance_interval_max
+                if guidance_interval_min != 0.0:
+                    interval_mask = interval_mask & (t_flow > guidance_interval_min)
+                scale = torch.where(
+                    interval_mask,
+                    torch.tensor(guidance_scale, device=latents.device, dtype=latents.dtype),
+                    torch.tensor(1.0, device=latents.device, dtype=latents.dtype),
+                )
+                x_pred = x_uncond + scale * (x_cond - x_uncond)
+
+            sigma = sigma.reshape(*([1] * (latents.ndim - 1)))
+            # JiT predicts x0; scheduler integrates in sigma space: dz/dsigma = -(x0 - z) / sigma.
+            model_output = -(x_pred - latents) / sigma
+            latents = self.scheduler.step(model_output, t, latents).prev_sample
 
         images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
         if output_type == "pt":

JiT-B-32/scheduler/scheduler_config.json

@@ -1,7 +1,6 @@
 {
-  "_class_name": "JiTScheduler",
+  "_class_name": "FlowMatchHeunDiscreteScheduler",
   "_diffusers_version": "0.36.0",
   "num_train_timesteps": 1000,
-  "t_eps": 0.05,
-  "solver": "heun"
+  "shift": 4.0
 }

JiT-B-32/transformer/jit_transformer_2d.py

@@ -68,38 +68,58 @@ class JiTRotaryEmbedding(nn.Module):
         num_cls_token=0,
     ):
         super().__init__()
-        if custom_freqs is not None:
-            freqs = custom_freqs
-        else:
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-
+        self.dim = dim
+        self.pt_seq_len = pt_seq_len
+        self.theta = theta
+        self.num_cls_token = num_cls_token
+        self.custom_freqs = custom_freqs
         if ft_seq_len is None:
             ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = freqs.repeat_interleave(2, dim=-1)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])  # [N_img, D]
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-
-            # prepend in-context cls token
-            _, D = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, D, dtype=cos_img.dtype)
-            sin_pad = torch.zeros(num_cls_token, D, dtype=sin_img.dtype)
-
-            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
-            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
+        self._cached_hw = None
+        cos, sin = self._build_freqs(ft_seq_len, ft_seq_len, device=torch.device("cpu"))
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+        self._cached_hw = (ft_seq_len, ft_seq_len)
+
+    def _build_freqs(self, height, width, device):
+        if self.custom_freqs is not None:
+            freqs = self.custom_freqs.to(device=device, dtype=torch.float32)
         else:
-            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
-            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
-
-    def forward(self, t):
+            freqs = 1.0 / (
+                self.theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)[: (self.dim // 2)] / self.dim)
+            )
+
+        t_h = torch.arange(height, device=device, dtype=torch.float32) / height * self.pt_seq_len
+        t_w = torch.arange(width, device=device, dtype=torch.float32) / width * self.pt_seq_len
+        freqs_h = torch.einsum("..., f -> ... f", t_h, freqs).repeat_interleave(2, dim=-1)
+        freqs_w = torch.einsum("..., f -> ... f", t_w, freqs).repeat_interleave(2, dim=-1)
+        freqs_2d = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
+        freqs_flat = freqs_2d.view(-1, freqs_2d.shape[-1])
+        cos_img = freqs_flat.cos()
+        sin_img = freqs_flat.sin()
+        if self.num_cls_token > 0:
+            _, dim_freq = cos_img.shape
+            cos_pad = torch.ones(self.num_cls_token, dim_freq, dtype=cos_img.dtype, device=device)
+            sin_pad = torch.zeros(self.num_cls_token, dim_freq, dtype=sin_img.dtype, device=device)
+            cos_img = torch.cat([cos_pad, cos_img], dim=0)
+            sin_img = torch.cat([sin_pad, sin_img], dim=0)
+        return cos_img, sin_img
+
+    def forward(self, t, height=None, width=None):
         # Applied on (batch, seq_len, heads, head_dim) tensors from attention.
         seq_len = t.shape[1]
+        if height is None or width is None:
+            image_tokens = seq_len - self.num_cls_token
+            size = int(image_tokens**0.5)
+            if size * size != image_tokens:
+                raise ValueError(
+                    f"Cannot infer square token grid from sequence length {seq_len} with {self.num_cls_token} class tokens."
+                )
+            height = size
+            width = size
+        if self._cached_hw != (height, width) or self.freqs_cos.device != t.device:
+            self.freqs_cos, self.freqs_sin = self._build_freqs(height, width, device=t.device)
+            self._cached_hw = (height, width)
         freqs_cos = self.freqs_cos[:seq_len].to(t.dtype)
         freqs_sin = self.freqs_sin[:seq_len].to(t.dtype)
 
@@ -195,7 +215,7 @@ class JiTAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
 
-    def forward(self, x, rope=None):
+    def forward(self, x, rope=None, grid_height=None, grid_width=None):
         B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
@@ -206,8 +226,8 @@ class JiTAttention(nn.Module):
         if rope is not None:
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
-            q = rope(q)
-            k = rope(k)
+            q = rope(q, height=grid_height, width=grid_width)
+            k = rope(k, height=grid_height, width=grid_width)
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
 
@@ -254,7 +274,7 @@ class JiTBlock(nn.Module):
         self.act = nn.SiLU()
         self.adaLN_modulation = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
 
-    def forward(self, x, c, feat_rope=None):
+    def forward(self, x, c, feat_rope=None, grid_height=None, grid_width=None):
         # Apply activation
         c = self.act(c)
 
@@ -263,7 +283,7 @@ class JiTBlock(nn.Module):
         # Attention block
         norm_x = self.norm1(x)
         modulated_x = modulate(norm_x, shift_msa, scale_msa)
-        attn_out = self.attn(modulated_x, rope=feat_rope)
+        attn_out = self.attn(modulated_x, rope=feat_rope, grid_height=grid_height, grid_width=grid_width)
         x = x + gate_msa.unsqueeze(1) * attn_out
 
         # MLP block
@@ -437,11 +457,30 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         self.act_final = nn.SiLU()
         self.adaLN_modulation_final = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
 
+    def _get_patch_grid(self, hidden_states):
+        height, width = hidden_states.shape[-2:]
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"Input size {(height, width)} must be divisible by patch_size={self.patch_size}."
+            )
+        return height // self.patch_size, width // self.patch_size
+
+    def _interpolate_pos_encoding(self, tokens, grid_height, grid_width):
+        num_tokens = grid_height * grid_width
+        if self.pos_embed.shape[1] == num_tokens:
+            return self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+        base_size = int(self.pos_embed.shape[1] ** 0.5)
+        pos_embed = self.pos_embed.reshape(1, base_size, base_size, self.hidden_size).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(grid_height, grid_width), mode="bicubic", align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(1, num_tokens, self.hidden_size)
+        return pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         class_labels: torch.LongTensor,
+        interpolate_pos_encoding: bool = True,
         return_dict: bool = True,
     ):
 
@@ -454,8 +493,19 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         c = t_emb + y_emb
 
         # Patch Embed
+        grid_height, grid_width = self._get_patch_grid(hidden_states)
         x = self.x_embedder(hidden_states)
-        x = x + self.pos_embed.to(x.dtype)
+        if interpolate_pos_encoding:
+            pos_embed = self._interpolate_pos_encoding(x, grid_height, grid_width)
+        else:
+            expected_tokens = grid_height * grid_width
+            if self.pos_embed.shape[1] != expected_tokens:
+                raise ValueError(
+                    f"pos_embed token count {self.pos_embed.shape[1]} does not match input token count {expected_tokens}. "
+                    "Enable interpolate_pos_encoding for dynamic resolutions."
+                )
+            pos_embed = self.pos_embed.to(device=x.device, dtype=x.dtype)
+        x = x + pos_embed
 
         # Blocks
         for i, block in enumerate(self.blocks):
@@ -467,15 +517,23 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
             rope = self.feat_rope if i < self.in_context_start else self.feat_rope_incontext
 
             if self.training and self.gradient_checkpointing:
+                def custom_forward(current_x, current_c):
+                    return block(
+                        current_x,
+                        current_c,
+                        feat_rope=rope,
+                        grid_height=grid_height,
+                        grid_width=grid_width,
+                    )
+
                 x = torch.utils.checkpoint.checkpoint(
-                    block,
+                    custom_forward,
                     x,
                     c,
-                    rope,
                     use_reentrant=False,
                 )
             else:
-                x = block(x, c, feat_rope=rope)
+                x = block(x, c, feat_rope=rope, grid_height=grid_height, grid_width=grid_width)
 
         # Slice off in-context tokens
         if self.in_context_len > 0:
@@ -489,10 +547,11 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         x = self.linear_final(x)
 
         # Unpatchify
-        h = w = int(x.shape[1] ** 0.5)
-        x = x.reshape(shape=(x.shape[0], h, w, self.patch_size, self.patch_size, self.out_channels))
+        x = x.reshape(shape=(x.shape[0], grid_height, grid_width, self.patch_size, self.patch_size, self.out_channels))
         x = torch.einsum("nhwpqc->nchpwq", x)
-        output = x.reshape(shape=(x.shape[0], self.out_channels, h * self.patch_size, w * self.patch_size))
+        output = x.reshape(
+            shape=(x.shape[0], self.out_channels, grid_height * self.patch_size, grid_width * self.patch_size)
+        )
 
         if not return_dict:
             return (output,)

JiT-H-16/model_index.json

@@ -5,11 +5,1013 @@
   ],
   "_diffusers_version": "0.36.0",
   "scheduler": [
-    "scheduling_jit",
-    "JiTScheduler"
+    "diffusers",
+    "FlowMatchHeunDiscreteScheduler"
   ],
   "transformer": [
     "jit_transformer_2d",
     "JiTTransformer2DModel"
-  ]
+  ],
+  "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"
+  }
 }

JiT-H-16/pipeline.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import importlib
 import json
 import sys
@@ -23,6 +21,7 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchHeunDiscreteScheduler, KarrasDiffusionSchedulers
 from diffusers.utils.torch_utils import randn_tensor
 
 
@@ -39,12 +38,10 @@ class JiTPipeline(DiffusionPipeline):
     Parameters:
         transformer ([`JiTTransformer2DModel`]):
             A class-conditioned `JiTTransformer2DModel` to denoise the images.
-        scheduler ([`JiTScheduler`]):
-            Manual JiT flow-matching scheduler (linear `t in [0, 1]`, Heun or Euler).
+        scheduler ([`KarrasDiffusionSchedulers`] or [`FlowMatchHeunDiscreteScheduler`]):
+            Diffusers scheduler interface for JiT generation (defaults to `FlowMatchHeunDiscreteScheduler(shift=4.0)`).
         id2label (`dict[int, str]`, *optional*):
             ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
-        id2label_cn (`dict[int, str]`, *optional*):
-            ImageNet class id to Chinese label mapping. Values may contain comma-separated synonyms.
     """
 
     model_cpu_offload_seq = "transformer"
@@ -71,7 +68,7 @@ class JiTPipeline(DiffusionPipeline):
 
             hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
             if subfolder:
-                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**", "labels/**"])
+                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
             cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
             variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
         else:
@@ -82,6 +79,7 @@ class JiTPipeline(DiffusionPipeline):
             if subfolder:
                 variant = variant / subfolder
 
+        id2label_override = kwargs.pop("id2label", None)
         model_kwargs = dict(kwargs)
         inserted: List[str] = []
 
@@ -103,19 +101,22 @@ class JiTPipeline(DiffusionPipeline):
 
         try:
             transformer = _load_component("transformer", "jit_transformer_2d", "JiTTransformer2DModel")
-            scheduler = _load_component("scheduler", "scheduling_jit", "JiTScheduler")
+            try:
+                scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained(str(variant), subfolder="scheduler")
+            except Exception:
+                scheduler = FlowMatchHeunDiscreteScheduler(shift=4.0)
 
             if transformer is None:
                 raise ValueError(f"No loadable transformer found under {variant}")
 
             variant_path = str(variant)
-            id2label, id2label_cn = cls._load_labels_for_variant(variant_path)
+            model_index_path = variant / "model_index.json"
+            id2label = id2label_override or cls._read_id2label_from_model_index(model_index_path)
 
             pipe = cls(
                 transformer=transformer,
                 scheduler=scheduler,
                 id2label=id2label,
-                id2label_cn=id2label_cn,
             )
             if variant_path and hasattr(pipe, "register_to_config"):
                 pipe.register_to_config(_name_or_path=variant_path)
@@ -128,58 +129,31 @@ class JiTPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler,
-        id2label: Optional[Dict[int, str]] = None,
-        id2label_cn: Optional[Dict[int, str]] = None,
+        scheduler: FlowMatchHeunDiscreteScheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
     ):
         super().__init__()
+        scheduler = scheduler or FlowMatchHeunDiscreteScheduler(shift=4.0)
         self.register_modules(transformer=transformer, scheduler=scheduler)
 
-        self._id2label = id2label or {}
-        self._id2label_cn = id2label_cn or {}
+        self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
-        self.labels_cn = self._build_label2id(self._id2label_cn)
-
-    def _ensure_labels_loaded(self) -> None:
-        if self._id2label or self._id2label_cn:
-            return
-        loaded_en, loaded_cn = self._load_labels_for_variant(getattr(self.config, "_name_or_path", None))
-        if loaded_en:
-            self._id2label = loaded_en
-            self.labels = self._build_label2id(self._id2label)
-        if loaded_cn:
-            self._id2label_cn = loaded_cn
-            self.labels_cn = self._build_label2id(self._id2label_cn)
 
     @staticmethod
-    def _labels_dir_for_variant(variant_path: Optional[str]) -> Optional[Path]:
-        if not variant_path:
-            return None
-        variant_dir = Path(variant_path).resolve()
-        labels_dir = variant_dir.parent / "labels"
-        return labels_dir if labels_dir.is_dir() else None
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
 
     @staticmethod
-    def _read_id2label(labels_dir: Path, lang: str = "en") -> Dict[int, str]:
-        filename = "id2label_en.json" if lang == "en" else "id2label_cn.json"
-        path = labels_dir / filename
-        if not path.exists():
-            raise FileNotFoundError(path)
-        raw = json.loads(path.read_text(encoding="utf-8"))
-        return {int(key): value for key, value in raw.items()}
-
-    @classmethod
-    def _load_labels_for_variant(
-        cls,
-        variant_path: Optional[str],
-    ) -> Tuple[Optional[Dict[int, str]], Optional[Dict[int, str]]]:
-        labels_dir = cls._labels_dir_for_variant(variant_path)
-        if labels_dir is None:
-            return None, None
-        try:
-            return cls._read_id2label(labels_dir, "en"), cls._read_id2label(labels_dir, "cn")
-        except FileNotFoundError:
-            return None, None
+    def _read_id2label_from_model_index(model_index_path: Path) -> Dict[int, str]:
+        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]:
@@ -194,35 +168,19 @@ class JiTPipeline(DiffusionPipeline):
     @property
     def id2label(self) -> Dict[int, str]:
         """ImageNet class id to English label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
         return self._id2label
 
-    @property
-    def id2label_cn(self) -> Dict[int, str]:
-        """ImageNet class id to Chinese label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
-        return self._id2label_cn
-
-    def get_label_ids(self, label: Union[str, List[str]], lang: str = "en") -> List[int]:
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
         r"""
         Map ImageNet label strings to class ids.
 
         Args:
             label (`str` or `list[str]`):
-                One or more label strings. Each string must match a synonym in `id2label` (English)
-                or `id2label_cn` (Chinese).
-            lang (`str`, *optional*, defaults to `"en"`):
-                `"en"` uses English synonyms; `"cn"` uses Chinese synonyms.
+                One or more English label strings. Each string must match a synonym in `id2label`.
         """
-        if lang not in ("en", "cn"):
-            raise ValueError(f"`lang` must be 'en' or 'cn', got {lang!r}.")
-
-        self._ensure_labels_loaded()
-        label2id = self.labels if lang == "en" else self.labels_cn
+        label2id = self.labels
         if not label2id:
-            raise ValueError(
-                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
-            )
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
 
         if isinstance(label, str):
             label = [label]
@@ -231,7 +189,7 @@ class JiTPipeline(DiffusionPipeline):
         if missing:
             preview = ", ".join(list(label2id.keys())[:8])
             raise ValueError(
-                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+                f"Unknown English label(s): {missing}. Example valid labels: {preview}, ..."
             )
         return [label2id[item] for item in label]
 
@@ -246,115 +204,10 @@ class JiTPipeline(DiffusionPipeline):
             return self.get_label_ids(class_labels)
 
         if class_labels and isinstance(class_labels[0], str):
-            self._ensure_labels_loaded()
-            if all(label in self.labels for label in class_labels):
-                return self.get_label_ids(class_labels, lang="en")
-            if all(label in self.labels_cn for label in class_labels):
-                return self.get_label_ids(class_labels, lang="cn")
-            raise ValueError(
-                "Could not resolve string `class_labels`. Use English synonyms from `pipe.labels` "
-                "or Chinese synonyms from `pipe.labels_cn`."
-            )
+            return self.get_label_ids(class_labels)
 
         return list(class_labels)
 
-    def _predict_velocity(
-        self,
-        z_value: torch.Tensor,
-        t: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-    ) -> torch.Tensor:
-        t = torch.as_tensor(t, device=z_value.device, dtype=z_value.dtype)
-        if do_classifier_free_guidance:
-            z_in = torch.cat([z_value, z_value], dim=0)
-            labels = torch.cat([class_labels, class_null], dim=0)
-        else:
-            z_in = z_value
-            labels = class_labels
-
-        t_batch = t.flatten().expand(z_in.shape[0])
-        x_pred = self.transformer(z_in, timestep=t_batch, class_labels=labels).sample
-        v = self.scheduler.velocity_from_prediction(z_in, x_pred, t)
-
-        if not do_classifier_free_guidance:
-            return v
-
-        v_cond, v_uncond = v.chunk(2, dim=0)
-        interval_mask = t < guidance_interval_max
-        if guidance_interval_min != 0.0:
-            interval_mask = interval_mask & (t > guidance_interval_min)
-        scale = torch.where(
-            interval_mask,
-            torch.tensor(guidance_scale, device=z_value.device, dtype=z_value.dtype),
-            torch.tensor(1.0, device=z_value.device, dtype=z_value.dtype),
-        )
-        return v_uncond + scale * (v_cond - v_uncond)
-
-    def _run_sampler(
-        self,
-        latents: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        num_inference_steps: int,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-        sampling_method: str,
-    ) -> torch.Tensor:
-        device = latents.device
-        self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampling_method)
-        timesteps = self.scheduler.timesteps
-
-        for i in self.progress_bar(range(num_inference_steps - 1)):
-            t = timesteps[i]
-            t_next = timesteps[i + 1]
-            v = self._predict_velocity(
-                latents,
-                t,
-                class_labels,
-                class_null,
-                do_classifier_free_guidance,
-                guidance_scale,
-                guidance_interval_min,
-                guidance_interval_max,
-            )
-
-            if sampling_method == "heun":
-                latents_euler = latents + (t_next - t) * v
-                v_next = self._predict_velocity(
-                    latents_euler,
-                    t_next,
-                    class_labels,
-                    class_null,
-                    do_classifier_free_guidance,
-                    guidance_scale,
-                    guidance_interval_min,
-                    guidance_interval_max,
-                )
-                latents = self.scheduler.step(v, t, latents, model_output_next=v_next).prev_sample
-            else:
-                latents = self.scheduler.step(v, t, latents).prev_sample
-
-        t = timesteps[-2]
-        t_next = timesteps[-1]
-        v = self._predict_velocity(
-            latents,
-            t,
-            class_labels,
-            class_null,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-        )
-        return latents + (t_next - t) * v
-
     @torch.inference_mode()
     def __call__(
         self,
@@ -363,10 +216,12 @@ class JiTPipeline(DiffusionPipeline):
         guidance_interval_min: float = 0.1,
         guidance_interval_max: float = 1.0,
         noise_scale: Optional[float] = None,
-        t_eps: Optional[float] = None,
-        sampling_method: Optional[str] = None,
+        t_eps: float = 5e-2,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         num_inference_steps: int = 50,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        interpolate_pos_encoding: bool = True,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
     ) -> Union[ImagePipelineOutput, Tuple]:
@@ -375,7 +230,7 @@ class JiTPipeline(DiffusionPipeline):
 
         Args:
             class_labels (`int`, `str`, `list[int]`, or `list[str]`):
-                ImageNet class indices or human-readable label strings (English or Chinese).
+                ImageNet class indices or human-readable English label strings.
             guidance_scale (`float`, *optional*):
                 Classifier-free guidance scale. CFG is active when `guidance_scale > 1.0`.
             guidance_interval_min (`float`, defaults to `0.1`):
@@ -384,10 +239,8 @@ class JiTPipeline(DiffusionPipeline):
                 Upper bound of the CFG interval in flow time.
             noise_scale (`float`, *optional*):
                 Initial Gaussian noise scale (`1.0` for 256px, `2.0` for 512px by default).
-            t_eps (`float`, *optional*):
-                Epsilon clamp for the `1 - t` denominator (scheduler config by default).
-            sampling_method (`str`, *optional*):
-                `"heun"` or `"euler"`. Defaults to the scheduler config (`heun`).
+            t_eps (`float`, defaults to `5e-2`):
+                Epsilon clamp for the `1 - t` denominator, matching JiT source defaults.
             generator (`torch.Generator`, *optional*):
                 RNG for reproducibility.
             num_inference_steps (`int`, defaults to `50`):
@@ -397,31 +250,34 @@ class JiTPipeline(DiffusionPipeline):
             return_dict (`bool`, *optional*, defaults to `True`):
                 Return [`ImagePipelineOutput`] if True.
         """
-        solver = sampling_method or self.scheduler.config.solver
-        if solver not in {"heun", "euler"}:
-            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
         if num_inference_steps < 2:
             raise ValueError("num_inference_steps must be >= 2.")
 
-        if t_eps is not None:
-            self.scheduler.register_to_config(t_eps=t_eps)
-
         class_label_ids = self._normalize_class_labels(class_labels)
         do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
 
         batch_size = len(class_label_ids)
         image_size = int(self.transformer.config.sample_size)
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
         channels = int(self.transformer.config.in_channels)
         null_class_val = int(self.transformer.config.num_classes)
 
         if guidance_scale is None:
             guidance_scale = 1.0
         if noise_scale is None:
-            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(image_size, 1.0)
+            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)
 
         latents = (
             randn_tensor(
-                shape=(batch_size, channels, image_size, image_size),
+                shape=(batch_size, channels, height, width),
                 generator=generator,
                 device=self._execution_device,
                 dtype=self.transformer.dtype,
@@ -433,17 +289,47 @@ class JiTPipeline(DiffusionPipeline):
         class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
         class_null = torch.full_like(class_labels_t, null_class_val)
 
-        latents = self._run_sampler(
-            latents,
-            class_labels_t,
-            class_null,
-            num_inference_steps,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-            solver,
-        )
+        if do_classifier_free_guidance:
+            class_labels_input = torch.cat([class_labels_t, class_null], dim=0)
+        else:
+            class_labels_input = class_labels_t
+
+        self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
+        for t in self.progress_bar(self.scheduler.timesteps):
+            step_index = self.scheduler.index_for_timestep(t, self.scheduler.timesteps)
+            sigma = self.scheduler.sigmas[step_index].to(device=latents.device, dtype=latents.dtype)
+            sigma = sigma.clamp_min(t_eps)
+            t_flow = (1.0 - sigma).clamp(0.0, 1.0)
+
+            if do_classifier_free_guidance:
+                latent_model_input = torch.cat([latents, latents], dim=0)
+            else:
+                latent_model_input = latents
+
+            timesteps = t_flow.flatten().expand(latent_model_input.shape[0])
+            x_pred = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=class_labels_input,
+                interpolate_pos_encoding=interpolate_pos_encoding,
+            ).sample
+
+            if do_classifier_free_guidance:
+                x_cond, x_uncond = x_pred.chunk(2, dim=0)
+                interval_mask = t_flow < guidance_interval_max
+                if guidance_interval_min != 0.0:
+                    interval_mask = interval_mask & (t_flow > guidance_interval_min)
+                scale = torch.where(
+                    interval_mask,
+                    torch.tensor(guidance_scale, device=latents.device, dtype=latents.dtype),
+                    torch.tensor(1.0, device=latents.device, dtype=latents.dtype),
+                )
+                x_pred = x_uncond + scale * (x_cond - x_uncond)
+
+            sigma = sigma.reshape(*([1] * (latents.ndim - 1)))
+            # JiT predicts x0; scheduler integrates in sigma space: dz/dsigma = -(x0 - z) / sigma.
+            model_output = -(x_pred - latents) / sigma
+            latents = self.scheduler.step(model_output, t, latents).prev_sample
 
         images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
         if output_type == "pt":

JiT-H-16/scheduler/scheduler_config.json

@@ -1,7 +1,6 @@
 {
-  "_class_name": "JiTScheduler",
+  "_class_name": "FlowMatchHeunDiscreteScheduler",
   "_diffusers_version": "0.36.0",
   "num_train_timesteps": 1000,
-  "t_eps": 0.05,
-  "solver": "heun"
+  "shift": 4.0
 }

JiT-H-16/transformer/jit_transformer_2d.py

@@ -68,38 +68,58 @@ class JiTRotaryEmbedding(nn.Module):
         num_cls_token=0,
     ):
         super().__init__()
-        if custom_freqs is not None:
-            freqs = custom_freqs
-        else:
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-
+        self.dim = dim
+        self.pt_seq_len = pt_seq_len
+        self.theta = theta
+        self.num_cls_token = num_cls_token
+        self.custom_freqs = custom_freqs
         if ft_seq_len is None:
             ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = freqs.repeat_interleave(2, dim=-1)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])  # [N_img, D]
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-
-            # prepend in-context cls token
-            _, D = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, D, dtype=cos_img.dtype)
-            sin_pad = torch.zeros(num_cls_token, D, dtype=sin_img.dtype)
-
-            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
-            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
+        self._cached_hw = None
+        cos, sin = self._build_freqs(ft_seq_len, ft_seq_len, device=torch.device("cpu"))
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+        self._cached_hw = (ft_seq_len, ft_seq_len)
+
+    def _build_freqs(self, height, width, device):
+        if self.custom_freqs is not None:
+            freqs = self.custom_freqs.to(device=device, dtype=torch.float32)
         else:
-            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
-            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
-
-    def forward(self, t):
+            freqs = 1.0 / (
+                self.theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)[: (self.dim // 2)] / self.dim)
+            )
+
+        t_h = torch.arange(height, device=device, dtype=torch.float32) / height * self.pt_seq_len
+        t_w = torch.arange(width, device=device, dtype=torch.float32) / width * self.pt_seq_len
+        freqs_h = torch.einsum("..., f -> ... f", t_h, freqs).repeat_interleave(2, dim=-1)
+        freqs_w = torch.einsum("..., f -> ... f", t_w, freqs).repeat_interleave(2, dim=-1)
+        freqs_2d = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
+        freqs_flat = freqs_2d.view(-1, freqs_2d.shape[-1])
+        cos_img = freqs_flat.cos()
+        sin_img = freqs_flat.sin()
+        if self.num_cls_token > 0:
+            _, dim_freq = cos_img.shape
+            cos_pad = torch.ones(self.num_cls_token, dim_freq, dtype=cos_img.dtype, device=device)
+            sin_pad = torch.zeros(self.num_cls_token, dim_freq, dtype=sin_img.dtype, device=device)
+            cos_img = torch.cat([cos_pad, cos_img], dim=0)
+            sin_img = torch.cat([sin_pad, sin_img], dim=0)
+        return cos_img, sin_img
+
+    def forward(self, t, height=None, width=None):
         # Applied on (batch, seq_len, heads, head_dim) tensors from attention.
         seq_len = t.shape[1]
+        if height is None or width is None:
+            image_tokens = seq_len - self.num_cls_token
+            size = int(image_tokens**0.5)
+            if size * size != image_tokens:
+                raise ValueError(
+                    f"Cannot infer square token grid from sequence length {seq_len} with {self.num_cls_token} class tokens."
+                )
+            height = size
+            width = size
+        if self._cached_hw != (height, width) or self.freqs_cos.device != t.device:
+            self.freqs_cos, self.freqs_sin = self._build_freqs(height, width, device=t.device)
+            self._cached_hw = (height, width)
         freqs_cos = self.freqs_cos[:seq_len].to(t.dtype)
         freqs_sin = self.freqs_sin[:seq_len].to(t.dtype)
 
@@ -195,7 +215,7 @@ class JiTAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
 
-    def forward(self, x, rope=None):
+    def forward(self, x, rope=None, grid_height=None, grid_width=None):
         B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
@@ -206,8 +226,8 @@ class JiTAttention(nn.Module):
         if rope is not None:
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
-            q = rope(q)
-            k = rope(k)
+            q = rope(q, height=grid_height, width=grid_width)
+            k = rope(k, height=grid_height, width=grid_width)
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
 
@@ -254,7 +274,7 @@ class JiTBlock(nn.Module):
         self.act = nn.SiLU()
         self.adaLN_modulation = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
 
-    def forward(self, x, c, feat_rope=None):
+    def forward(self, x, c, feat_rope=None, grid_height=None, grid_width=None):
         # Apply activation
         c = self.act(c)
 
@@ -263,7 +283,7 @@ class JiTBlock(nn.Module):
         # Attention block
         norm_x = self.norm1(x)
         modulated_x = modulate(norm_x, shift_msa, scale_msa)
-        attn_out = self.attn(modulated_x, rope=feat_rope)
+        attn_out = self.attn(modulated_x, rope=feat_rope, grid_height=grid_height, grid_width=grid_width)
         x = x + gate_msa.unsqueeze(1) * attn_out
 
         # MLP block
@@ -437,11 +457,30 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         self.act_final = nn.SiLU()
         self.adaLN_modulation_final = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
 
+    def _get_patch_grid(self, hidden_states):
+        height, width = hidden_states.shape[-2:]
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"Input size {(height, width)} must be divisible by patch_size={self.patch_size}."
+            )
+        return height // self.patch_size, width // self.patch_size
+
+    def _interpolate_pos_encoding(self, tokens, grid_height, grid_width):
+        num_tokens = grid_height * grid_width
+        if self.pos_embed.shape[1] == num_tokens:
+            return self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+        base_size = int(self.pos_embed.shape[1] ** 0.5)
+        pos_embed = self.pos_embed.reshape(1, base_size, base_size, self.hidden_size).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(grid_height, grid_width), mode="bicubic", align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(1, num_tokens, self.hidden_size)
+        return pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         class_labels: torch.LongTensor,
+        interpolate_pos_encoding: bool = True,
         return_dict: bool = True,
     ):
 
@@ -454,8 +493,19 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         c = t_emb + y_emb
 
         # Patch Embed
+        grid_height, grid_width = self._get_patch_grid(hidden_states)
         x = self.x_embedder(hidden_states)
-        x = x + self.pos_embed.to(x.dtype)
+        if interpolate_pos_encoding:
+            pos_embed = self._interpolate_pos_encoding(x, grid_height, grid_width)
+        else:
+            expected_tokens = grid_height * grid_width
+            if self.pos_embed.shape[1] != expected_tokens:
+                raise ValueError(
+                    f"pos_embed token count {self.pos_embed.shape[1]} does not match input token count {expected_tokens}. "
+                    "Enable interpolate_pos_encoding for dynamic resolutions."
+                )
+            pos_embed = self.pos_embed.to(device=x.device, dtype=x.dtype)
+        x = x + pos_embed
 
         # Blocks
         for i, block in enumerate(self.blocks):
@@ -467,15 +517,23 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
             rope = self.feat_rope if i < self.in_context_start else self.feat_rope_incontext
 
             if self.training and self.gradient_checkpointing:
+                def custom_forward(current_x, current_c):
+                    return block(
+                        current_x,
+                        current_c,
+                        feat_rope=rope,
+                        grid_height=grid_height,
+                        grid_width=grid_width,
+                    )
+
                 x = torch.utils.checkpoint.checkpoint(
-                    block,
+                    custom_forward,
                     x,
                     c,
-                    rope,
                     use_reentrant=False,
                 )
             else:
-                x = block(x, c, feat_rope=rope)
+                x = block(x, c, feat_rope=rope, grid_height=grid_height, grid_width=grid_width)
 
         # Slice off in-context tokens
         if self.in_context_len > 0:
@@ -489,10 +547,11 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         x = self.linear_final(x)
 
         # Unpatchify
-        h = w = int(x.shape[1] ** 0.5)
-        x = x.reshape(shape=(x.shape[0], h, w, self.patch_size, self.patch_size, self.out_channels))
+        x = x.reshape(shape=(x.shape[0], grid_height, grid_width, self.patch_size, self.patch_size, self.out_channels))
         x = torch.einsum("nhwpqc->nchpwq", x)
-        output = x.reshape(shape=(x.shape[0], self.out_channels, h * self.patch_size, w * self.patch_size))
+        output = x.reshape(
+            shape=(x.shape[0], self.out_channels, grid_height * self.patch_size, grid_width * self.patch_size)
+        )
 
         if not return_dict:
             return (output,)

JiT-H-32/model_index.json

@@ -5,11 +5,1013 @@
   ],
   "_diffusers_version": "0.36.0",
   "scheduler": [
-    "scheduling_jit",
-    "JiTScheduler"
+    "diffusers",
+    "FlowMatchHeunDiscreteScheduler"
   ],
   "transformer": [
     "jit_transformer_2d",
     "JiTTransformer2DModel"
-  ]
+  ],
+  "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"
+  }
 }

JiT-H-32/pipeline.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import importlib
 import json
 import sys
@@ -23,6 +21,7 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchHeunDiscreteScheduler, KarrasDiffusionSchedulers
 from diffusers.utils.torch_utils import randn_tensor
 
 
@@ -39,12 +38,10 @@ class JiTPipeline(DiffusionPipeline):
     Parameters:
         transformer ([`JiTTransformer2DModel`]):
             A class-conditioned `JiTTransformer2DModel` to denoise the images.
-        scheduler ([`JiTScheduler`]):
-            Manual JiT flow-matching scheduler (linear `t in [0, 1]`, Heun or Euler).
+        scheduler ([`KarrasDiffusionSchedulers`] or [`FlowMatchHeunDiscreteScheduler`]):
+            Diffusers scheduler interface for JiT generation (defaults to `FlowMatchHeunDiscreteScheduler(shift=4.0)`).
         id2label (`dict[int, str]`, *optional*):
             ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
-        id2label_cn (`dict[int, str]`, *optional*):
-            ImageNet class id to Chinese label mapping. Values may contain comma-separated synonyms.
     """
 
     model_cpu_offload_seq = "transformer"
@@ -71,7 +68,7 @@ class JiTPipeline(DiffusionPipeline):
 
             hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
             if subfolder:
-                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**", "labels/**"])
+                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
             cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
             variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
         else:
@@ -82,6 +79,7 @@ class JiTPipeline(DiffusionPipeline):
             if subfolder:
                 variant = variant / subfolder
 
+        id2label_override = kwargs.pop("id2label", None)
         model_kwargs = dict(kwargs)
         inserted: List[str] = []
 
@@ -103,19 +101,22 @@ class JiTPipeline(DiffusionPipeline):
 
         try:
             transformer = _load_component("transformer", "jit_transformer_2d", "JiTTransformer2DModel")
-            scheduler = _load_component("scheduler", "scheduling_jit", "JiTScheduler")
+            try:
+                scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained(str(variant), subfolder="scheduler")
+            except Exception:
+                scheduler = FlowMatchHeunDiscreteScheduler(shift=4.0)
 
             if transformer is None:
                 raise ValueError(f"No loadable transformer found under {variant}")
 
             variant_path = str(variant)
-            id2label, id2label_cn = cls._load_labels_for_variant(variant_path)
+            model_index_path = variant / "model_index.json"
+            id2label = id2label_override or cls._read_id2label_from_model_index(model_index_path)
 
             pipe = cls(
                 transformer=transformer,
                 scheduler=scheduler,
                 id2label=id2label,
-                id2label_cn=id2label_cn,
             )
             if variant_path and hasattr(pipe, "register_to_config"):
                 pipe.register_to_config(_name_or_path=variant_path)
@@ -128,58 +129,31 @@ class JiTPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler,
-        id2label: Optional[Dict[int, str]] = None,
-        id2label_cn: Optional[Dict[int, str]] = None,
+        scheduler: FlowMatchHeunDiscreteScheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
     ):
         super().__init__()
+        scheduler = scheduler or FlowMatchHeunDiscreteScheduler(shift=4.0)
         self.register_modules(transformer=transformer, scheduler=scheduler)
 
-        self._id2label = id2label or {}
-        self._id2label_cn = id2label_cn or {}
+        self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
-        self.labels_cn = self._build_label2id(self._id2label_cn)
-
-    def _ensure_labels_loaded(self) -> None:
-        if self._id2label or self._id2label_cn:
-            return
-        loaded_en, loaded_cn = self._load_labels_for_variant(getattr(self.config, "_name_or_path", None))
-        if loaded_en:
-            self._id2label = loaded_en
-            self.labels = self._build_label2id(self._id2label)
-        if loaded_cn:
-            self._id2label_cn = loaded_cn
-            self.labels_cn = self._build_label2id(self._id2label_cn)
 
     @staticmethod
-    def _labels_dir_for_variant(variant_path: Optional[str]) -> Optional[Path]:
-        if not variant_path:
-            return None
-        variant_dir = Path(variant_path).resolve()
-        labels_dir = variant_dir.parent / "labels"
-        return labels_dir if labels_dir.is_dir() else None
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
 
     @staticmethod
-    def _read_id2label(labels_dir: Path, lang: str = "en") -> Dict[int, str]:
-        filename = "id2label_en.json" if lang == "en" else "id2label_cn.json"
-        path = labels_dir / filename
-        if not path.exists():
-            raise FileNotFoundError(path)
-        raw = json.loads(path.read_text(encoding="utf-8"))
-        return {int(key): value for key, value in raw.items()}
-
-    @classmethod
-    def _load_labels_for_variant(
-        cls,
-        variant_path: Optional[str],
-    ) -> Tuple[Optional[Dict[int, str]], Optional[Dict[int, str]]]:
-        labels_dir = cls._labels_dir_for_variant(variant_path)
-        if labels_dir is None:
-            return None, None
-        try:
-            return cls._read_id2label(labels_dir, "en"), cls._read_id2label(labels_dir, "cn")
-        except FileNotFoundError:
-            return None, None
+    def _read_id2label_from_model_index(model_index_path: Path) -> Dict[int, str]:
+        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]:
@@ -194,35 +168,19 @@ class JiTPipeline(DiffusionPipeline):
     @property
     def id2label(self) -> Dict[int, str]:
         """ImageNet class id to English label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
         return self._id2label
 
-    @property
-    def id2label_cn(self) -> Dict[int, str]:
-        """ImageNet class id to Chinese label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
-        return self._id2label_cn
-
-    def get_label_ids(self, label: Union[str, List[str]], lang: str = "en") -> List[int]:
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
         r"""
         Map ImageNet label strings to class ids.
 
         Args:
             label (`str` or `list[str]`):
-                One or more label strings. Each string must match a synonym in `id2label` (English)
-                or `id2label_cn` (Chinese).
-            lang (`str`, *optional*, defaults to `"en"`):
-                `"en"` uses English synonyms; `"cn"` uses Chinese synonyms.
+                One or more English label strings. Each string must match a synonym in `id2label`.
         """
-        if lang not in ("en", "cn"):
-            raise ValueError(f"`lang` must be 'en' or 'cn', got {lang!r}.")
-
-        self._ensure_labels_loaded()
-        label2id = self.labels if lang == "en" else self.labels_cn
+        label2id = self.labels
         if not label2id:
-            raise ValueError(
-                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
-            )
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
 
         if isinstance(label, str):
             label = [label]
@@ -231,7 +189,7 @@ class JiTPipeline(DiffusionPipeline):
         if missing:
             preview = ", ".join(list(label2id.keys())[:8])
             raise ValueError(
-                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+                f"Unknown English label(s): {missing}. Example valid labels: {preview}, ..."
             )
         return [label2id[item] for item in label]
 
@@ -246,115 +204,10 @@ class JiTPipeline(DiffusionPipeline):
             return self.get_label_ids(class_labels)
 
         if class_labels and isinstance(class_labels[0], str):
-            self._ensure_labels_loaded()
-            if all(label in self.labels for label in class_labels):
-                return self.get_label_ids(class_labels, lang="en")
-            if all(label in self.labels_cn for label in class_labels):
-                return self.get_label_ids(class_labels, lang="cn")
-            raise ValueError(
-                "Could not resolve string `class_labels`. Use English synonyms from `pipe.labels` "
-                "or Chinese synonyms from `pipe.labels_cn`."
-            )
+            return self.get_label_ids(class_labels)
 
         return list(class_labels)
 
-    def _predict_velocity(
-        self,
-        z_value: torch.Tensor,
-        t: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-    ) -> torch.Tensor:
-        t = torch.as_tensor(t, device=z_value.device, dtype=z_value.dtype)
-        if do_classifier_free_guidance:
-            z_in = torch.cat([z_value, z_value], dim=0)
-            labels = torch.cat([class_labels, class_null], dim=0)
-        else:
-            z_in = z_value
-            labels = class_labels
-
-        t_batch = t.flatten().expand(z_in.shape[0])
-        x_pred = self.transformer(z_in, timestep=t_batch, class_labels=labels).sample
-        v = self.scheduler.velocity_from_prediction(z_in, x_pred, t)
-
-        if not do_classifier_free_guidance:
-            return v
-
-        v_cond, v_uncond = v.chunk(2, dim=0)
-        interval_mask = t < guidance_interval_max
-        if guidance_interval_min != 0.0:
-            interval_mask = interval_mask & (t > guidance_interval_min)
-        scale = torch.where(
-            interval_mask,
-            torch.tensor(guidance_scale, device=z_value.device, dtype=z_value.dtype),
-            torch.tensor(1.0, device=z_value.device, dtype=z_value.dtype),
-        )
-        return v_uncond + scale * (v_cond - v_uncond)
-
-    def _run_sampler(
-        self,
-        latents: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        num_inference_steps: int,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-        sampling_method: str,
-    ) -> torch.Tensor:
-        device = latents.device
-        self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampling_method)
-        timesteps = self.scheduler.timesteps
-
-        for i in self.progress_bar(range(num_inference_steps - 1)):
-            t = timesteps[i]
-            t_next = timesteps[i + 1]
-            v = self._predict_velocity(
-                latents,
-                t,
-                class_labels,
-                class_null,
-                do_classifier_free_guidance,
-                guidance_scale,
-                guidance_interval_min,
-                guidance_interval_max,
-            )
-
-            if sampling_method == "heun":
-                latents_euler = latents + (t_next - t) * v
-                v_next = self._predict_velocity(
-                    latents_euler,
-                    t_next,
-                    class_labels,
-                    class_null,
-                    do_classifier_free_guidance,
-                    guidance_scale,
-                    guidance_interval_min,
-                    guidance_interval_max,
-                )
-                latents = self.scheduler.step(v, t, latents, model_output_next=v_next).prev_sample
-            else:
-                latents = self.scheduler.step(v, t, latents).prev_sample
-
-        t = timesteps[-2]
-        t_next = timesteps[-1]
-        v = self._predict_velocity(
-            latents,
-            t,
-            class_labels,
-            class_null,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-        )
-        return latents + (t_next - t) * v
-
     @torch.inference_mode()
     def __call__(
         self,
@@ -363,10 +216,12 @@ class JiTPipeline(DiffusionPipeline):
         guidance_interval_min: float = 0.1,
         guidance_interval_max: float = 1.0,
         noise_scale: Optional[float] = None,
-        t_eps: Optional[float] = None,
-        sampling_method: Optional[str] = None,
+        t_eps: float = 5e-2,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         num_inference_steps: int = 50,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        interpolate_pos_encoding: bool = True,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
     ) -> Union[ImagePipelineOutput, Tuple]:
@@ -375,7 +230,7 @@ class JiTPipeline(DiffusionPipeline):
 
         Args:
             class_labels (`int`, `str`, `list[int]`, or `list[str]`):
-                ImageNet class indices or human-readable label strings (English or Chinese).
+                ImageNet class indices or human-readable English label strings.
             guidance_scale (`float`, *optional*):
                 Classifier-free guidance scale. CFG is active when `guidance_scale > 1.0`.
             guidance_interval_min (`float`, defaults to `0.1`):
@@ -384,10 +239,8 @@ class JiTPipeline(DiffusionPipeline):
                 Upper bound of the CFG interval in flow time.
             noise_scale (`float`, *optional*):
                 Initial Gaussian noise scale (`1.0` for 256px, `2.0` for 512px by default).
-            t_eps (`float`, *optional*):
-                Epsilon clamp for the `1 - t` denominator (scheduler config by default).
-            sampling_method (`str`, *optional*):
-                `"heun"` or `"euler"`. Defaults to the scheduler config (`heun`).
+            t_eps (`float`, defaults to `5e-2`):
+                Epsilon clamp for the `1 - t` denominator, matching JiT source defaults.
             generator (`torch.Generator`, *optional*):
                 RNG for reproducibility.
             num_inference_steps (`int`, defaults to `50`):
@@ -397,31 +250,34 @@ class JiTPipeline(DiffusionPipeline):
             return_dict (`bool`, *optional*, defaults to `True`):
                 Return [`ImagePipelineOutput`] if True.
         """
-        solver = sampling_method or self.scheduler.config.solver
-        if solver not in {"heun", "euler"}:
-            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
         if num_inference_steps < 2:
             raise ValueError("num_inference_steps must be >= 2.")
 
-        if t_eps is not None:
-            self.scheduler.register_to_config(t_eps=t_eps)
-
         class_label_ids = self._normalize_class_labels(class_labels)
         do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
 
         batch_size = len(class_label_ids)
         image_size = int(self.transformer.config.sample_size)
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
         channels = int(self.transformer.config.in_channels)
         null_class_val = int(self.transformer.config.num_classes)
 
         if guidance_scale is None:
             guidance_scale = 1.0
         if noise_scale is None:
-            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(image_size, 1.0)
+            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)
 
         latents = (
             randn_tensor(
-                shape=(batch_size, channels, image_size, image_size),
+                shape=(batch_size, channels, height, width),
                 generator=generator,
                 device=self._execution_device,
                 dtype=self.transformer.dtype,
@@ -433,17 +289,47 @@ class JiTPipeline(DiffusionPipeline):
         class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
         class_null = torch.full_like(class_labels_t, null_class_val)
 
-        latents = self._run_sampler(
-            latents,
-            class_labels_t,
-            class_null,
-            num_inference_steps,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-            solver,
-        )
+        if do_classifier_free_guidance:
+            class_labels_input = torch.cat([class_labels_t, class_null], dim=0)
+        else:
+            class_labels_input = class_labels_t
+
+        self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
+        for t in self.progress_bar(self.scheduler.timesteps):
+            step_index = self.scheduler.index_for_timestep(t, self.scheduler.timesteps)
+            sigma = self.scheduler.sigmas[step_index].to(device=latents.device, dtype=latents.dtype)
+            sigma = sigma.clamp_min(t_eps)
+            t_flow = (1.0 - sigma).clamp(0.0, 1.0)
+
+            if do_classifier_free_guidance:
+                latent_model_input = torch.cat([latents, latents], dim=0)
+            else:
+                latent_model_input = latents
+
+            timesteps = t_flow.flatten().expand(latent_model_input.shape[0])
+            x_pred = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=class_labels_input,
+                interpolate_pos_encoding=interpolate_pos_encoding,
+            ).sample
+
+            if do_classifier_free_guidance:
+                x_cond, x_uncond = x_pred.chunk(2, dim=0)
+                interval_mask = t_flow < guidance_interval_max
+                if guidance_interval_min != 0.0:
+                    interval_mask = interval_mask & (t_flow > guidance_interval_min)
+                scale = torch.where(
+                    interval_mask,
+                    torch.tensor(guidance_scale, device=latents.device, dtype=latents.dtype),
+                    torch.tensor(1.0, device=latents.device, dtype=latents.dtype),
+                )
+                x_pred = x_uncond + scale * (x_cond - x_uncond)
+
+            sigma = sigma.reshape(*([1] * (latents.ndim - 1)))
+            # JiT predicts x0; scheduler integrates in sigma space: dz/dsigma = -(x0 - z) / sigma.
+            model_output = -(x_pred - latents) / sigma
+            latents = self.scheduler.step(model_output, t, latents).prev_sample
 
         images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
         if output_type == "pt":

JiT-H-32/scheduler/scheduler_config.json

@@ -1,7 +1,6 @@
 {
-  "_class_name": "JiTScheduler",
+  "_class_name": "FlowMatchHeunDiscreteScheduler",
   "_diffusers_version": "0.36.0",
   "num_train_timesteps": 1000,
-  "t_eps": 0.05,
-  "solver": "heun"
+  "shift": 4.0
 }

JiT-H-32/transformer/jit_transformer_2d.py

@@ -68,38 +68,58 @@ class JiTRotaryEmbedding(nn.Module):
         num_cls_token=0,
     ):
         super().__init__()
-        if custom_freqs is not None:
-            freqs = custom_freqs
-        else:
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-
+        self.dim = dim
+        self.pt_seq_len = pt_seq_len
+        self.theta = theta
+        self.num_cls_token = num_cls_token
+        self.custom_freqs = custom_freqs
         if ft_seq_len is None:
             ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = freqs.repeat_interleave(2, dim=-1)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])  # [N_img, D]
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-
-            # prepend in-context cls token
-            _, D = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, D, dtype=cos_img.dtype)
-            sin_pad = torch.zeros(num_cls_token, D, dtype=sin_img.dtype)
-
-            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
-            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
+        self._cached_hw = None
+        cos, sin = self._build_freqs(ft_seq_len, ft_seq_len, device=torch.device("cpu"))
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+        self._cached_hw = (ft_seq_len, ft_seq_len)
+
+    def _build_freqs(self, height, width, device):
+        if self.custom_freqs is not None:
+            freqs = self.custom_freqs.to(device=device, dtype=torch.float32)
         else:
-            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
-            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
-
-    def forward(self, t):
+            freqs = 1.0 / (
+                self.theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)[: (self.dim // 2)] / self.dim)
+            )
+
+        t_h = torch.arange(height, device=device, dtype=torch.float32) / height * self.pt_seq_len
+        t_w = torch.arange(width, device=device, dtype=torch.float32) / width * self.pt_seq_len
+        freqs_h = torch.einsum("..., f -> ... f", t_h, freqs).repeat_interleave(2, dim=-1)
+        freqs_w = torch.einsum("..., f -> ... f", t_w, freqs).repeat_interleave(2, dim=-1)
+        freqs_2d = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
+        freqs_flat = freqs_2d.view(-1, freqs_2d.shape[-1])
+        cos_img = freqs_flat.cos()
+        sin_img = freqs_flat.sin()
+        if self.num_cls_token > 0:
+            _, dim_freq = cos_img.shape
+            cos_pad = torch.ones(self.num_cls_token, dim_freq, dtype=cos_img.dtype, device=device)
+            sin_pad = torch.zeros(self.num_cls_token, dim_freq, dtype=sin_img.dtype, device=device)
+            cos_img = torch.cat([cos_pad, cos_img], dim=0)
+            sin_img = torch.cat([sin_pad, sin_img], dim=0)
+        return cos_img, sin_img
+
+    def forward(self, t, height=None, width=None):
         # Applied on (batch, seq_len, heads, head_dim) tensors from attention.
         seq_len = t.shape[1]
+        if height is None or width is None:
+            image_tokens = seq_len - self.num_cls_token
+            size = int(image_tokens**0.5)
+            if size * size != image_tokens:
+                raise ValueError(
+                    f"Cannot infer square token grid from sequence length {seq_len} with {self.num_cls_token} class tokens."
+                )
+            height = size
+            width = size
+        if self._cached_hw != (height, width) or self.freqs_cos.device != t.device:
+            self.freqs_cos, self.freqs_sin = self._build_freqs(height, width, device=t.device)
+            self._cached_hw = (height, width)
         freqs_cos = self.freqs_cos[:seq_len].to(t.dtype)
         freqs_sin = self.freqs_sin[:seq_len].to(t.dtype)
 
@@ -195,7 +215,7 @@ class JiTAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
 
-    def forward(self, x, rope=None):
+    def forward(self, x, rope=None, grid_height=None, grid_width=None):
         B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
@@ -206,8 +226,8 @@ class JiTAttention(nn.Module):
         if rope is not None:
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
-            q = rope(q)
-            k = rope(k)
+            q = rope(q, height=grid_height, width=grid_width)
+            k = rope(k, height=grid_height, width=grid_width)
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
 
@@ -254,7 +274,7 @@ class JiTBlock(nn.Module):
         self.act = nn.SiLU()
         self.adaLN_modulation = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
 
-    def forward(self, x, c, feat_rope=None):
+    def forward(self, x, c, feat_rope=None, grid_height=None, grid_width=None):
         # Apply activation
         c = self.act(c)
 
@@ -263,7 +283,7 @@ class JiTBlock(nn.Module):
         # Attention block
         norm_x = self.norm1(x)
         modulated_x = modulate(norm_x, shift_msa, scale_msa)
-        attn_out = self.attn(modulated_x, rope=feat_rope)
+        attn_out = self.attn(modulated_x, rope=feat_rope, grid_height=grid_height, grid_width=grid_width)
         x = x + gate_msa.unsqueeze(1) * attn_out
 
         # MLP block
@@ -437,11 +457,30 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         self.act_final = nn.SiLU()
         self.adaLN_modulation_final = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
 
+    def _get_patch_grid(self, hidden_states):
+        height, width = hidden_states.shape[-2:]
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"Input size {(height, width)} must be divisible by patch_size={self.patch_size}."
+            )
+        return height // self.patch_size, width // self.patch_size
+
+    def _interpolate_pos_encoding(self, tokens, grid_height, grid_width):
+        num_tokens = grid_height * grid_width
+        if self.pos_embed.shape[1] == num_tokens:
+            return self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+        base_size = int(self.pos_embed.shape[1] ** 0.5)
+        pos_embed = self.pos_embed.reshape(1, base_size, base_size, self.hidden_size).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(grid_height, grid_width), mode="bicubic", align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(1, num_tokens, self.hidden_size)
+        return pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         class_labels: torch.LongTensor,
+        interpolate_pos_encoding: bool = True,
         return_dict: bool = True,
     ):
 
@@ -454,8 +493,19 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         c = t_emb + y_emb
 
         # Patch Embed
+        grid_height, grid_width = self._get_patch_grid(hidden_states)
         x = self.x_embedder(hidden_states)
-        x = x + self.pos_embed.to(x.dtype)
+        if interpolate_pos_encoding:
+            pos_embed = self._interpolate_pos_encoding(x, grid_height, grid_width)
+        else:
+            expected_tokens = grid_height * grid_width
+            if self.pos_embed.shape[1] != expected_tokens:
+                raise ValueError(
+                    f"pos_embed token count {self.pos_embed.shape[1]} does not match input token count {expected_tokens}. "
+                    "Enable interpolate_pos_encoding for dynamic resolutions."
+                )
+            pos_embed = self.pos_embed.to(device=x.device, dtype=x.dtype)
+        x = x + pos_embed
 
         # Blocks
         for i, block in enumerate(self.blocks):
@@ -467,15 +517,23 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
             rope = self.feat_rope if i < self.in_context_start else self.feat_rope_incontext
 
             if self.training and self.gradient_checkpointing:
+                def custom_forward(current_x, current_c):
+                    return block(
+                        current_x,
+                        current_c,
+                        feat_rope=rope,
+                        grid_height=grid_height,
+                        grid_width=grid_width,
+                    )
+
                 x = torch.utils.checkpoint.checkpoint(
-                    block,
+                    custom_forward,
                     x,
                     c,
-                    rope,
                     use_reentrant=False,
                 )
             else:
-                x = block(x, c, feat_rope=rope)
+                x = block(x, c, feat_rope=rope, grid_height=grid_height, grid_width=grid_width)
 
         # Slice off in-context tokens
         if self.in_context_len > 0:
@@ -489,10 +547,11 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         x = self.linear_final(x)
 
         # Unpatchify
-        h = w = int(x.shape[1] ** 0.5)
-        x = x.reshape(shape=(x.shape[0], h, w, self.patch_size, self.patch_size, self.out_channels))
+        x = x.reshape(shape=(x.shape[0], grid_height, grid_width, self.patch_size, self.patch_size, self.out_channels))
         x = torch.einsum("nhwpqc->nchpwq", x)
-        output = x.reshape(shape=(x.shape[0], self.out_channels, h * self.patch_size, w * self.patch_size))
+        output = x.reshape(
+            shape=(x.shape[0], self.out_channels, grid_height * self.patch_size, grid_width * self.patch_size)
+        )
 
         if not return_dict:
             return (output,)

JiT-L-16/model_index.json

@@ -5,11 +5,1013 @@
   ],
   "_diffusers_version": "0.36.0",
   "scheduler": [
-    "scheduling_jit",
-    "JiTScheduler"
+    "diffusers",
+    "FlowMatchHeunDiscreteScheduler"
   ],
   "transformer": [
     "jit_transformer_2d",
     "JiTTransformer2DModel"
-  ]
+  ],
+  "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"
+  }
 }

JiT-L-16/pipeline.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import importlib
 import json
 import sys
@@ -23,6 +21,7 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchHeunDiscreteScheduler, KarrasDiffusionSchedulers
 from diffusers.utils.torch_utils import randn_tensor
 
 
@@ -39,12 +38,10 @@ class JiTPipeline(DiffusionPipeline):
     Parameters:
         transformer ([`JiTTransformer2DModel`]):
             A class-conditioned `JiTTransformer2DModel` to denoise the images.
-        scheduler ([`JiTScheduler`]):
-            Manual JiT flow-matching scheduler (linear `t in [0, 1]`, Heun or Euler).
+        scheduler ([`KarrasDiffusionSchedulers`] or [`FlowMatchHeunDiscreteScheduler`]):
+            Diffusers scheduler interface for JiT generation (defaults to `FlowMatchHeunDiscreteScheduler(shift=4.0)`).
         id2label (`dict[int, str]`, *optional*):
             ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
-        id2label_cn (`dict[int, str]`, *optional*):
-            ImageNet class id to Chinese label mapping. Values may contain comma-separated synonyms.
     """
 
     model_cpu_offload_seq = "transformer"
@@ -71,7 +68,7 @@ class JiTPipeline(DiffusionPipeline):
 
             hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
             if subfolder:
-                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**", "labels/**"])
+                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
             cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
             variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
         else:
@@ -82,6 +79,7 @@ class JiTPipeline(DiffusionPipeline):
             if subfolder:
                 variant = variant / subfolder
 
+        id2label_override = kwargs.pop("id2label", None)
         model_kwargs = dict(kwargs)
         inserted: List[str] = []
 
@@ -103,19 +101,22 @@ class JiTPipeline(DiffusionPipeline):
 
         try:
             transformer = _load_component("transformer", "jit_transformer_2d", "JiTTransformer2DModel")
-            scheduler = _load_component("scheduler", "scheduling_jit", "JiTScheduler")
+            try:
+                scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained(str(variant), subfolder="scheduler")
+            except Exception:
+                scheduler = FlowMatchHeunDiscreteScheduler(shift=4.0)
 
             if transformer is None:
                 raise ValueError(f"No loadable transformer found under {variant}")
 
             variant_path = str(variant)
-            id2label, id2label_cn = cls._load_labels_for_variant(variant_path)
+            model_index_path = variant / "model_index.json"
+            id2label = id2label_override or cls._read_id2label_from_model_index(model_index_path)
 
             pipe = cls(
                 transformer=transformer,
                 scheduler=scheduler,
                 id2label=id2label,
-                id2label_cn=id2label_cn,
             )
             if variant_path and hasattr(pipe, "register_to_config"):
                 pipe.register_to_config(_name_or_path=variant_path)
@@ -128,58 +129,31 @@ class JiTPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler,
-        id2label: Optional[Dict[int, str]] = None,
-        id2label_cn: Optional[Dict[int, str]] = None,
+        scheduler: FlowMatchHeunDiscreteScheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
     ):
         super().__init__()
+        scheduler = scheduler or FlowMatchHeunDiscreteScheduler(shift=4.0)
         self.register_modules(transformer=transformer, scheduler=scheduler)
 
-        self._id2label = id2label or {}
-        self._id2label_cn = id2label_cn or {}
+        self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
-        self.labels_cn = self._build_label2id(self._id2label_cn)
-
-    def _ensure_labels_loaded(self) -> None:
-        if self._id2label or self._id2label_cn:
-            return
-        loaded_en, loaded_cn = self._load_labels_for_variant(getattr(self.config, "_name_or_path", None))
-        if loaded_en:
-            self._id2label = loaded_en
-            self.labels = self._build_label2id(self._id2label)
-        if loaded_cn:
-            self._id2label_cn = loaded_cn
-            self.labels_cn = self._build_label2id(self._id2label_cn)
 
     @staticmethod
-    def _labels_dir_for_variant(variant_path: Optional[str]) -> Optional[Path]:
-        if not variant_path:
-            return None
-        variant_dir = Path(variant_path).resolve()
-        labels_dir = variant_dir.parent / "labels"
-        return labels_dir if labels_dir.is_dir() else None
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
 
     @staticmethod
-    def _read_id2label(labels_dir: Path, lang: str = "en") -> Dict[int, str]:
-        filename = "id2label_en.json" if lang == "en" else "id2label_cn.json"
-        path = labels_dir / filename
-        if not path.exists():
-            raise FileNotFoundError(path)
-        raw = json.loads(path.read_text(encoding="utf-8"))
-        return {int(key): value for key, value in raw.items()}
-
-    @classmethod
-    def _load_labels_for_variant(
-        cls,
-        variant_path: Optional[str],
-    ) -> Tuple[Optional[Dict[int, str]], Optional[Dict[int, str]]]:
-        labels_dir = cls._labels_dir_for_variant(variant_path)
-        if labels_dir is None:
-            return None, None
-        try:
-            return cls._read_id2label(labels_dir, "en"), cls._read_id2label(labels_dir, "cn")
-        except FileNotFoundError:
-            return None, None
+    def _read_id2label_from_model_index(model_index_path: Path) -> Dict[int, str]:
+        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]:
@@ -194,35 +168,19 @@ class JiTPipeline(DiffusionPipeline):
     @property
     def id2label(self) -> Dict[int, str]:
         """ImageNet class id to English label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
         return self._id2label
 
-    @property
-    def id2label_cn(self) -> Dict[int, str]:
-        """ImageNet class id to Chinese label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
-        return self._id2label_cn
-
-    def get_label_ids(self, label: Union[str, List[str]], lang: str = "en") -> List[int]:
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
         r"""
         Map ImageNet label strings to class ids.
 
         Args:
             label (`str` or `list[str]`):
-                One or more label strings. Each string must match a synonym in `id2label` (English)
-                or `id2label_cn` (Chinese).
-            lang (`str`, *optional*, defaults to `"en"`):
-                `"en"` uses English synonyms; `"cn"` uses Chinese synonyms.
+                One or more English label strings. Each string must match a synonym in `id2label`.
         """
-        if lang not in ("en", "cn"):
-            raise ValueError(f"`lang` must be 'en' or 'cn', got {lang!r}.")
-
-        self._ensure_labels_loaded()
-        label2id = self.labels if lang == "en" else self.labels_cn
+        label2id = self.labels
         if not label2id:
-            raise ValueError(
-                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
-            )
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
 
         if isinstance(label, str):
             label = [label]
@@ -231,7 +189,7 @@ class JiTPipeline(DiffusionPipeline):
         if missing:
             preview = ", ".join(list(label2id.keys())[:8])
             raise ValueError(
-                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+                f"Unknown English label(s): {missing}. Example valid labels: {preview}, ..."
             )
         return [label2id[item] for item in label]
 
@@ -246,115 +204,10 @@ class JiTPipeline(DiffusionPipeline):
             return self.get_label_ids(class_labels)
 
         if class_labels and isinstance(class_labels[0], str):
-            self._ensure_labels_loaded()
-            if all(label in self.labels for label in class_labels):
-                return self.get_label_ids(class_labels, lang="en")
-            if all(label in self.labels_cn for label in class_labels):
-                return self.get_label_ids(class_labels, lang="cn")
-            raise ValueError(
-                "Could not resolve string `class_labels`. Use English synonyms from `pipe.labels` "
-                "or Chinese synonyms from `pipe.labels_cn`."
-            )
+            return self.get_label_ids(class_labels)
 
         return list(class_labels)
 
-    def _predict_velocity(
-        self,
-        z_value: torch.Tensor,
-        t: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-    ) -> torch.Tensor:
-        t = torch.as_tensor(t, device=z_value.device, dtype=z_value.dtype)
-        if do_classifier_free_guidance:
-            z_in = torch.cat([z_value, z_value], dim=0)
-            labels = torch.cat([class_labels, class_null], dim=0)
-        else:
-            z_in = z_value
-            labels = class_labels
-
-        t_batch = t.flatten().expand(z_in.shape[0])
-        x_pred = self.transformer(z_in, timestep=t_batch, class_labels=labels).sample
-        v = self.scheduler.velocity_from_prediction(z_in, x_pred, t)
-
-        if not do_classifier_free_guidance:
-            return v
-
-        v_cond, v_uncond = v.chunk(2, dim=0)
-        interval_mask = t < guidance_interval_max
-        if guidance_interval_min != 0.0:
-            interval_mask = interval_mask & (t > guidance_interval_min)
-        scale = torch.where(
-            interval_mask,
-            torch.tensor(guidance_scale, device=z_value.device, dtype=z_value.dtype),
-            torch.tensor(1.0, device=z_value.device, dtype=z_value.dtype),
-        )
-        return v_uncond + scale * (v_cond - v_uncond)
-
-    def _run_sampler(
-        self,
-        latents: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        num_inference_steps: int,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-        sampling_method: str,
-    ) -> torch.Tensor:
-        device = latents.device
-        self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampling_method)
-        timesteps = self.scheduler.timesteps
-
-        for i in self.progress_bar(range(num_inference_steps - 1)):
-            t = timesteps[i]
-            t_next = timesteps[i + 1]
-            v = self._predict_velocity(
-                latents,
-                t,
-                class_labels,
-                class_null,
-                do_classifier_free_guidance,
-                guidance_scale,
-                guidance_interval_min,
-                guidance_interval_max,
-            )
-
-            if sampling_method == "heun":
-                latents_euler = latents + (t_next - t) * v
-                v_next = self._predict_velocity(
-                    latents_euler,
-                    t_next,
-                    class_labels,
-                    class_null,
-                    do_classifier_free_guidance,
-                    guidance_scale,
-                    guidance_interval_min,
-                    guidance_interval_max,
-                )
-                latents = self.scheduler.step(v, t, latents, model_output_next=v_next).prev_sample
-            else:
-                latents = self.scheduler.step(v, t, latents).prev_sample
-
-        t = timesteps[-2]
-        t_next = timesteps[-1]
-        v = self._predict_velocity(
-            latents,
-            t,
-            class_labels,
-            class_null,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-        )
-        return latents + (t_next - t) * v
-
     @torch.inference_mode()
     def __call__(
         self,
@@ -363,10 +216,12 @@ class JiTPipeline(DiffusionPipeline):
         guidance_interval_min: float = 0.1,
         guidance_interval_max: float = 1.0,
         noise_scale: Optional[float] = None,
-        t_eps: Optional[float] = None,
-        sampling_method: Optional[str] = None,
+        t_eps: float = 5e-2,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         num_inference_steps: int = 50,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        interpolate_pos_encoding: bool = True,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
     ) -> Union[ImagePipelineOutput, Tuple]:
@@ -375,7 +230,7 @@ class JiTPipeline(DiffusionPipeline):
 
         Args:
             class_labels (`int`, `str`, `list[int]`, or `list[str]`):
-                ImageNet class indices or human-readable label strings (English or Chinese).
+                ImageNet class indices or human-readable English label strings.
             guidance_scale (`float`, *optional*):
                 Classifier-free guidance scale. CFG is active when `guidance_scale > 1.0`.
             guidance_interval_min (`float`, defaults to `0.1`):
@@ -384,10 +239,8 @@ class JiTPipeline(DiffusionPipeline):
                 Upper bound of the CFG interval in flow time.
             noise_scale (`float`, *optional*):
                 Initial Gaussian noise scale (`1.0` for 256px, `2.0` for 512px by default).
-            t_eps (`float`, *optional*):
-                Epsilon clamp for the `1 - t` denominator (scheduler config by default).
-            sampling_method (`str`, *optional*):
-                `"heun"` or `"euler"`. Defaults to the scheduler config (`heun`).
+            t_eps (`float`, defaults to `5e-2`):
+                Epsilon clamp for the `1 - t` denominator, matching JiT source defaults.
             generator (`torch.Generator`, *optional*):
                 RNG for reproducibility.
             num_inference_steps (`int`, defaults to `50`):
@@ -397,31 +250,34 @@ class JiTPipeline(DiffusionPipeline):
             return_dict (`bool`, *optional*, defaults to `True`):
                 Return [`ImagePipelineOutput`] if True.
         """
-        solver = sampling_method or self.scheduler.config.solver
-        if solver not in {"heun", "euler"}:
-            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
         if num_inference_steps < 2:
             raise ValueError("num_inference_steps must be >= 2.")
 
-        if t_eps is not None:
-            self.scheduler.register_to_config(t_eps=t_eps)
-
         class_label_ids = self._normalize_class_labels(class_labels)
         do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
 
         batch_size = len(class_label_ids)
         image_size = int(self.transformer.config.sample_size)
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
         channels = int(self.transformer.config.in_channels)
         null_class_val = int(self.transformer.config.num_classes)
 
         if guidance_scale is None:
             guidance_scale = 1.0
         if noise_scale is None:
-            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(image_size, 1.0)
+            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)
 
         latents = (
             randn_tensor(
-                shape=(batch_size, channels, image_size, image_size),
+                shape=(batch_size, channels, height, width),
                 generator=generator,
                 device=self._execution_device,
                 dtype=self.transformer.dtype,
@@ -433,17 +289,47 @@ class JiTPipeline(DiffusionPipeline):
         class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
         class_null = torch.full_like(class_labels_t, null_class_val)
 
-        latents = self._run_sampler(
-            latents,
-            class_labels_t,
-            class_null,
-            num_inference_steps,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-            solver,
-        )
+        if do_classifier_free_guidance:
+            class_labels_input = torch.cat([class_labels_t, class_null], dim=0)
+        else:
+            class_labels_input = class_labels_t
+
+        self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
+        for t in self.progress_bar(self.scheduler.timesteps):
+            step_index = self.scheduler.index_for_timestep(t, self.scheduler.timesteps)
+            sigma = self.scheduler.sigmas[step_index].to(device=latents.device, dtype=latents.dtype)
+            sigma = sigma.clamp_min(t_eps)
+            t_flow = (1.0 - sigma).clamp(0.0, 1.0)
+
+            if do_classifier_free_guidance:
+                latent_model_input = torch.cat([latents, latents], dim=0)
+            else:
+                latent_model_input = latents
+
+            timesteps = t_flow.flatten().expand(latent_model_input.shape[0])
+            x_pred = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=class_labels_input,
+                interpolate_pos_encoding=interpolate_pos_encoding,
+            ).sample
+
+            if do_classifier_free_guidance:
+                x_cond, x_uncond = x_pred.chunk(2, dim=0)
+                interval_mask = t_flow < guidance_interval_max
+                if guidance_interval_min != 0.0:
+                    interval_mask = interval_mask & (t_flow > guidance_interval_min)
+                scale = torch.where(
+                    interval_mask,
+                    torch.tensor(guidance_scale, device=latents.device, dtype=latents.dtype),
+                    torch.tensor(1.0, device=latents.device, dtype=latents.dtype),
+                )
+                x_pred = x_uncond + scale * (x_cond - x_uncond)
+
+            sigma = sigma.reshape(*([1] * (latents.ndim - 1)))
+            # JiT predicts x0; scheduler integrates in sigma space: dz/dsigma = -(x0 - z) / sigma.
+            model_output = -(x_pred - latents) / sigma
+            latents = self.scheduler.step(model_output, t, latents).prev_sample
 
         images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
         if output_type == "pt":

JiT-L-16/scheduler/scheduler_config.json

@@ -1,7 +1,6 @@
 {
-  "_class_name": "JiTScheduler",
+  "_class_name": "FlowMatchHeunDiscreteScheduler",
   "_diffusers_version": "0.36.0",
   "num_train_timesteps": 1000,
-  "t_eps": 0.05,
-  "solver": "heun"
+  "shift": 4.0
 }

JiT-L-16/transformer/jit_transformer_2d.py

@@ -68,38 +68,58 @@ class JiTRotaryEmbedding(nn.Module):
         num_cls_token=0,
     ):
         super().__init__()
-        if custom_freqs is not None:
-            freqs = custom_freqs
-        else:
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-
+        self.dim = dim
+        self.pt_seq_len = pt_seq_len
+        self.theta = theta
+        self.num_cls_token = num_cls_token
+        self.custom_freqs = custom_freqs
         if ft_seq_len is None:
             ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = freqs.repeat_interleave(2, dim=-1)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])  # [N_img, D]
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-
-            # prepend in-context cls token
-            _, D = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, D, dtype=cos_img.dtype)
-            sin_pad = torch.zeros(num_cls_token, D, dtype=sin_img.dtype)
-
-            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
-            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
+        self._cached_hw = None
+        cos, sin = self._build_freqs(ft_seq_len, ft_seq_len, device=torch.device("cpu"))
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+        self._cached_hw = (ft_seq_len, ft_seq_len)
+
+    def _build_freqs(self, height, width, device):
+        if self.custom_freqs is not None:
+            freqs = self.custom_freqs.to(device=device, dtype=torch.float32)
         else:
-            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
-            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
-
-    def forward(self, t):
+            freqs = 1.0 / (
+                self.theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)[: (self.dim // 2)] / self.dim)
+            )
+
+        t_h = torch.arange(height, device=device, dtype=torch.float32) / height * self.pt_seq_len
+        t_w = torch.arange(width, device=device, dtype=torch.float32) / width * self.pt_seq_len
+        freqs_h = torch.einsum("..., f -> ... f", t_h, freqs).repeat_interleave(2, dim=-1)
+        freqs_w = torch.einsum("..., f -> ... f", t_w, freqs).repeat_interleave(2, dim=-1)
+        freqs_2d = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
+        freqs_flat = freqs_2d.view(-1, freqs_2d.shape[-1])
+        cos_img = freqs_flat.cos()
+        sin_img = freqs_flat.sin()
+        if self.num_cls_token > 0:
+            _, dim_freq = cos_img.shape
+            cos_pad = torch.ones(self.num_cls_token, dim_freq, dtype=cos_img.dtype, device=device)
+            sin_pad = torch.zeros(self.num_cls_token, dim_freq, dtype=sin_img.dtype, device=device)
+            cos_img = torch.cat([cos_pad, cos_img], dim=0)
+            sin_img = torch.cat([sin_pad, sin_img], dim=0)
+        return cos_img, sin_img
+
+    def forward(self, t, height=None, width=None):
         # Applied on (batch, seq_len, heads, head_dim) tensors from attention.
         seq_len = t.shape[1]
+        if height is None or width is None:
+            image_tokens = seq_len - self.num_cls_token
+            size = int(image_tokens**0.5)
+            if size * size != image_tokens:
+                raise ValueError(
+                    f"Cannot infer square token grid from sequence length {seq_len} with {self.num_cls_token} class tokens."
+                )
+            height = size
+            width = size
+        if self._cached_hw != (height, width) or self.freqs_cos.device != t.device:
+            self.freqs_cos, self.freqs_sin = self._build_freqs(height, width, device=t.device)
+            self._cached_hw = (height, width)
         freqs_cos = self.freqs_cos[:seq_len].to(t.dtype)
         freqs_sin = self.freqs_sin[:seq_len].to(t.dtype)
 
@@ -195,7 +215,7 @@ class JiTAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
 
-    def forward(self, x, rope=None):
+    def forward(self, x, rope=None, grid_height=None, grid_width=None):
         B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
@@ -206,8 +226,8 @@ class JiTAttention(nn.Module):
         if rope is not None:
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
-            q = rope(q)
-            k = rope(k)
+            q = rope(q, height=grid_height, width=grid_width)
+            k = rope(k, height=grid_height, width=grid_width)
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
 
@@ -254,7 +274,7 @@ class JiTBlock(nn.Module):
         self.act = nn.SiLU()
         self.adaLN_modulation = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
 
-    def forward(self, x, c, feat_rope=None):
+    def forward(self, x, c, feat_rope=None, grid_height=None, grid_width=None):
         # Apply activation
         c = self.act(c)
 
@@ -263,7 +283,7 @@ class JiTBlock(nn.Module):
         # Attention block
         norm_x = self.norm1(x)
         modulated_x = modulate(norm_x, shift_msa, scale_msa)
-        attn_out = self.attn(modulated_x, rope=feat_rope)
+        attn_out = self.attn(modulated_x, rope=feat_rope, grid_height=grid_height, grid_width=grid_width)
         x = x + gate_msa.unsqueeze(1) * attn_out
 
         # MLP block
@@ -437,11 +457,30 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         self.act_final = nn.SiLU()
         self.adaLN_modulation_final = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
 
+    def _get_patch_grid(self, hidden_states):
+        height, width = hidden_states.shape[-2:]
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"Input size {(height, width)} must be divisible by patch_size={self.patch_size}."
+            )
+        return height // self.patch_size, width // self.patch_size
+
+    def _interpolate_pos_encoding(self, tokens, grid_height, grid_width):
+        num_tokens = grid_height * grid_width
+        if self.pos_embed.shape[1] == num_tokens:
+            return self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+        base_size = int(self.pos_embed.shape[1] ** 0.5)
+        pos_embed = self.pos_embed.reshape(1, base_size, base_size, self.hidden_size).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(grid_height, grid_width), mode="bicubic", align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(1, num_tokens, self.hidden_size)
+        return pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         class_labels: torch.LongTensor,
+        interpolate_pos_encoding: bool = True,
         return_dict: bool = True,
     ):
 
@@ -454,8 +493,19 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         c = t_emb + y_emb
 
         # Patch Embed
+        grid_height, grid_width = self._get_patch_grid(hidden_states)
         x = self.x_embedder(hidden_states)
-        x = x + self.pos_embed.to(x.dtype)
+        if interpolate_pos_encoding:
+            pos_embed = self._interpolate_pos_encoding(x, grid_height, grid_width)
+        else:
+            expected_tokens = grid_height * grid_width
+            if self.pos_embed.shape[1] != expected_tokens:
+                raise ValueError(
+                    f"pos_embed token count {self.pos_embed.shape[1]} does not match input token count {expected_tokens}. "
+                    "Enable interpolate_pos_encoding for dynamic resolutions."
+                )
+            pos_embed = self.pos_embed.to(device=x.device, dtype=x.dtype)
+        x = x + pos_embed
 
         # Blocks
         for i, block in enumerate(self.blocks):
@@ -467,15 +517,23 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
             rope = self.feat_rope if i < self.in_context_start else self.feat_rope_incontext
 
             if self.training and self.gradient_checkpointing:
+                def custom_forward(current_x, current_c):
+                    return block(
+                        current_x,
+                        current_c,
+                        feat_rope=rope,
+                        grid_height=grid_height,
+                        grid_width=grid_width,
+                    )
+
                 x = torch.utils.checkpoint.checkpoint(
-                    block,
+                    custom_forward,
                     x,
                     c,
-                    rope,
                     use_reentrant=False,
                 )
             else:
-                x = block(x, c, feat_rope=rope)
+                x = block(x, c, feat_rope=rope, grid_height=grid_height, grid_width=grid_width)
 
         # Slice off in-context tokens
         if self.in_context_len > 0:
@@ -489,10 +547,11 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         x = self.linear_final(x)
 
         # Unpatchify
-        h = w = int(x.shape[1] ** 0.5)
-        x = x.reshape(shape=(x.shape[0], h, w, self.patch_size, self.patch_size, self.out_channels))
+        x = x.reshape(shape=(x.shape[0], grid_height, grid_width, self.patch_size, self.patch_size, self.out_channels))
         x = torch.einsum("nhwpqc->nchpwq", x)
-        output = x.reshape(shape=(x.shape[0], self.out_channels, h * self.patch_size, w * self.patch_size))
+        output = x.reshape(
+            shape=(x.shape[0], self.out_channels, grid_height * self.patch_size, grid_width * self.patch_size)
+        )
 
         if not return_dict:
             return (output,)

JiT-L-32/model_index.json

@@ -5,11 +5,1013 @@
   ],
   "_diffusers_version": "0.36.0",
   "scheduler": [
-    "scheduling_jit",
-    "JiTScheduler"
+    "diffusers",
+    "FlowMatchHeunDiscreteScheduler"
   ],
   "transformer": [
     "jit_transformer_2d",
     "JiTTransformer2DModel"
-  ]
+  ],
+  "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"
+  }
 }

JiT-L-32/pipeline.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import importlib
 import json
 import sys
@@ -23,6 +21,7 @@ from typing import Dict, List, Optional, Tuple, Union
 import torch
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import FlowMatchHeunDiscreteScheduler, KarrasDiffusionSchedulers
 from diffusers.utils.torch_utils import randn_tensor
 
 
@@ -39,12 +38,10 @@ class JiTPipeline(DiffusionPipeline):
     Parameters:
         transformer ([`JiTTransformer2DModel`]):
             A class-conditioned `JiTTransformer2DModel` to denoise the images.
-        scheduler ([`JiTScheduler`]):
-            Manual JiT flow-matching scheduler (linear `t in [0, 1]`, Heun or Euler).
+        scheduler ([`KarrasDiffusionSchedulers`] or [`FlowMatchHeunDiscreteScheduler`]):
+            Diffusers scheduler interface for JiT generation (defaults to `FlowMatchHeunDiscreteScheduler(shift=4.0)`).
         id2label (`dict[int, str]`, *optional*):
             ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
-        id2label_cn (`dict[int, str]`, *optional*):
-            ImageNet class id to Chinese label mapping. Values may contain comma-separated synonyms.
     """
 
     model_cpu_offload_seq = "transformer"
@@ -71,7 +68,7 @@ class JiTPipeline(DiffusionPipeline):
 
             hub_kwargs = dict(kwargs.pop("hub_kwargs", {}))
             if subfolder:
-                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**", "labels/**"])
+                hub_kwargs.setdefault("allow_patterns", [f"{subfolder}/**"])
             cache_dir = snapshot_download(pretrained_model_name_or_path, **hub_kwargs)
             variant = Path(cache_dir) / subfolder if subfolder else Path(cache_dir)
         else:
@@ -82,6 +79,7 @@ class JiTPipeline(DiffusionPipeline):
             if subfolder:
                 variant = variant / subfolder
 
+        id2label_override = kwargs.pop("id2label", None)
         model_kwargs = dict(kwargs)
         inserted: List[str] = []
 
@@ -103,19 +101,22 @@ class JiTPipeline(DiffusionPipeline):
 
         try:
             transformer = _load_component("transformer", "jit_transformer_2d", "JiTTransformer2DModel")
-            scheduler = _load_component("scheduler", "scheduling_jit", "JiTScheduler")
+            try:
+                scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained(str(variant), subfolder="scheduler")
+            except Exception:
+                scheduler = FlowMatchHeunDiscreteScheduler(shift=4.0)
 
             if transformer is None:
                 raise ValueError(f"No loadable transformer found under {variant}")
 
             variant_path = str(variant)
-            id2label, id2label_cn = cls._load_labels_for_variant(variant_path)
+            model_index_path = variant / "model_index.json"
+            id2label = id2label_override or cls._read_id2label_from_model_index(model_index_path)
 
             pipe = cls(
                 transformer=transformer,
                 scheduler=scheduler,
                 id2label=id2label,
-                id2label_cn=id2label_cn,
             )
             if variant_path and hasattr(pipe, "register_to_config"):
                 pipe.register_to_config(_name_or_path=variant_path)
@@ -128,58 +129,31 @@ class JiTPipeline(DiffusionPipeline):
     def __init__(
         self,
         transformer,
-        scheduler,
-        id2label: Optional[Dict[int, str]] = None,
-        id2label_cn: Optional[Dict[int, str]] = None,
+        scheduler: FlowMatchHeunDiscreteScheduler,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
     ):
         super().__init__()
+        scheduler = scheduler or FlowMatchHeunDiscreteScheduler(shift=4.0)
         self.register_modules(transformer=transformer, scheduler=scheduler)
 
-        self._id2label = id2label or {}
-        self._id2label_cn = id2label_cn or {}
+        self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
-        self.labels_cn = self._build_label2id(self._id2label_cn)
-
-    def _ensure_labels_loaded(self) -> None:
-        if self._id2label or self._id2label_cn:
-            return
-        loaded_en, loaded_cn = self._load_labels_for_variant(getattr(self.config, "_name_or_path", None))
-        if loaded_en:
-            self._id2label = loaded_en
-            self.labels = self._build_label2id(self._id2label)
-        if loaded_cn:
-            self._id2label_cn = loaded_cn
-            self.labels_cn = self._build_label2id(self._id2label_cn)
 
     @staticmethod
-    def _labels_dir_for_variant(variant_path: Optional[str]) -> Optional[Path]:
-        if not variant_path:
-            return None
-        variant_dir = Path(variant_path).resolve()
-        labels_dir = variant_dir.parent / "labels"
-        return labels_dir if labels_dir.is_dir() else None
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            return {}
+        return {int(key): value for key, value in id2label.items()}
 
     @staticmethod
-    def _read_id2label(labels_dir: Path, lang: str = "en") -> Dict[int, str]:
-        filename = "id2label_en.json" if lang == "en" else "id2label_cn.json"
-        path = labels_dir / filename
-        if not path.exists():
-            raise FileNotFoundError(path)
-        raw = json.loads(path.read_text(encoding="utf-8"))
-        return {int(key): value for key, value in raw.items()}
-
-    @classmethod
-    def _load_labels_for_variant(
-        cls,
-        variant_path: Optional[str],
-    ) -> Tuple[Optional[Dict[int, str]], Optional[Dict[int, str]]]:
-        labels_dir = cls._labels_dir_for_variant(variant_path)
-        if labels_dir is None:
-            return None, None
-        try:
-            return cls._read_id2label(labels_dir, "en"), cls._read_id2label(labels_dir, "cn")
-        except FileNotFoundError:
-            return None, None
+    def _read_id2label_from_model_index(model_index_path: Path) -> Dict[int, str]:
+        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]:
@@ -194,35 +168,19 @@ class JiTPipeline(DiffusionPipeline):
     @property
     def id2label(self) -> Dict[int, str]:
         """ImageNet class id to English label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
         return self._id2label
 
-    @property
-    def id2label_cn(self) -> Dict[int, str]:
-        """ImageNet class id to Chinese label string (comma-separated synonyms)."""
-        self._ensure_labels_loaded()
-        return self._id2label_cn
-
-    def get_label_ids(self, label: Union[str, List[str]], lang: str = "en") -> List[int]:
+    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
         r"""
         Map ImageNet label strings to class ids.
 
         Args:
             label (`str` or `list[str]`):
-                One or more label strings. Each string must match a synonym in `id2label` (English)
-                or `id2label_cn` (Chinese).
-            lang (`str`, *optional*, defaults to `"en"`):
-                `"en"` uses English synonyms; `"cn"` uses Chinese synonyms.
+                One or more English label strings. Each string must match a synonym in `id2label`.
         """
-        if lang not in ("en", "cn"):
-            raise ValueError(f"`lang` must be 'en' or 'cn', got {lang!r}.")
-
-        self._ensure_labels_loaded()
-        label2id = self.labels if lang == "en" else self.labels_cn
+        label2id = self.labels
         if not label2id:
-            raise ValueError(
-                f"No {lang} labels loaded. Ensure `labels/id2label_{lang}.json` exists next to the variant folder."
-            )
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
 
         if isinstance(label, str):
             label = [label]
@@ -231,7 +189,7 @@ class JiTPipeline(DiffusionPipeline):
         if missing:
             preview = ", ".join(list(label2id.keys())[:8])
             raise ValueError(
-                f"Unknown label(s) for lang={lang!r}: {missing}. Example valid labels: {preview}, ..."
+                f"Unknown English label(s): {missing}. Example valid labels: {preview}, ..."
             )
         return [label2id[item] for item in label]
 
@@ -246,115 +204,10 @@ class JiTPipeline(DiffusionPipeline):
             return self.get_label_ids(class_labels)
 
         if class_labels and isinstance(class_labels[0], str):
-            self._ensure_labels_loaded()
-            if all(label in self.labels for label in class_labels):
-                return self.get_label_ids(class_labels, lang="en")
-            if all(label in self.labels_cn for label in class_labels):
-                return self.get_label_ids(class_labels, lang="cn")
-            raise ValueError(
-                "Could not resolve string `class_labels`. Use English synonyms from `pipe.labels` "
-                "or Chinese synonyms from `pipe.labels_cn`."
-            )
+            return self.get_label_ids(class_labels)
 
         return list(class_labels)
 
-    def _predict_velocity(
-        self,
-        z_value: torch.Tensor,
-        t: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-    ) -> torch.Tensor:
-        t = torch.as_tensor(t, device=z_value.device, dtype=z_value.dtype)
-        if do_classifier_free_guidance:
-            z_in = torch.cat([z_value, z_value], dim=0)
-            labels = torch.cat([class_labels, class_null], dim=0)
-        else:
-            z_in = z_value
-            labels = class_labels
-
-        t_batch = t.flatten().expand(z_in.shape[0])
-        x_pred = self.transformer(z_in, timestep=t_batch, class_labels=labels).sample
-        v = self.scheduler.velocity_from_prediction(z_in, x_pred, t)
-
-        if not do_classifier_free_guidance:
-            return v
-
-        v_cond, v_uncond = v.chunk(2, dim=0)
-        interval_mask = t < guidance_interval_max
-        if guidance_interval_min != 0.0:
-            interval_mask = interval_mask & (t > guidance_interval_min)
-        scale = torch.where(
-            interval_mask,
-            torch.tensor(guidance_scale, device=z_value.device, dtype=z_value.dtype),
-            torch.tensor(1.0, device=z_value.device, dtype=z_value.dtype),
-        )
-        return v_uncond + scale * (v_cond - v_uncond)
-
-    def _run_sampler(
-        self,
-        latents: torch.Tensor,
-        class_labels: torch.Tensor,
-        class_null: torch.Tensor,
-        num_inference_steps: int,
-        do_classifier_free_guidance: bool,
-        guidance_scale: float,
-        guidance_interval_min: float,
-        guidance_interval_max: float,
-        sampling_method: str,
-    ) -> torch.Tensor:
-        device = latents.device
-        self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampling_method)
-        timesteps = self.scheduler.timesteps
-
-        for i in self.progress_bar(range(num_inference_steps - 1)):
-            t = timesteps[i]
-            t_next = timesteps[i + 1]
-            v = self._predict_velocity(
-                latents,
-                t,
-                class_labels,
-                class_null,
-                do_classifier_free_guidance,
-                guidance_scale,
-                guidance_interval_min,
-                guidance_interval_max,
-            )
-
-            if sampling_method == "heun":
-                latents_euler = latents + (t_next - t) * v
-                v_next = self._predict_velocity(
-                    latents_euler,
-                    t_next,
-                    class_labels,
-                    class_null,
-                    do_classifier_free_guidance,
-                    guidance_scale,
-                    guidance_interval_min,
-                    guidance_interval_max,
-                )
-                latents = self.scheduler.step(v, t, latents, model_output_next=v_next).prev_sample
-            else:
-                latents = self.scheduler.step(v, t, latents).prev_sample
-
-        t = timesteps[-2]
-        t_next = timesteps[-1]
-        v = self._predict_velocity(
-            latents,
-            t,
-            class_labels,
-            class_null,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-        )
-        return latents + (t_next - t) * v
-
     @torch.inference_mode()
     def __call__(
         self,
@@ -363,10 +216,12 @@ class JiTPipeline(DiffusionPipeline):
         guidance_interval_min: float = 0.1,
         guidance_interval_max: float = 1.0,
         noise_scale: Optional[float] = None,
-        t_eps: Optional[float] = None,
-        sampling_method: Optional[str] = None,
+        t_eps: float = 5e-2,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         num_inference_steps: int = 50,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        interpolate_pos_encoding: bool = True,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
     ) -> Union[ImagePipelineOutput, Tuple]:
@@ -375,7 +230,7 @@ class JiTPipeline(DiffusionPipeline):
 
         Args:
             class_labels (`int`, `str`, `list[int]`, or `list[str]`):
-                ImageNet class indices or human-readable label strings (English or Chinese).
+                ImageNet class indices or human-readable English label strings.
             guidance_scale (`float`, *optional*):
                 Classifier-free guidance scale. CFG is active when `guidance_scale > 1.0`.
             guidance_interval_min (`float`, defaults to `0.1`):
@@ -384,10 +239,8 @@ class JiTPipeline(DiffusionPipeline):
                 Upper bound of the CFG interval in flow time.
             noise_scale (`float`, *optional*):
                 Initial Gaussian noise scale (`1.0` for 256px, `2.0` for 512px by default).
-            t_eps (`float`, *optional*):
-                Epsilon clamp for the `1 - t` denominator (scheduler config by default).
-            sampling_method (`str`, *optional*):
-                `"heun"` or `"euler"`. Defaults to the scheduler config (`heun`).
+            t_eps (`float`, defaults to `5e-2`):
+                Epsilon clamp for the `1 - t` denominator, matching JiT source defaults.
             generator (`torch.Generator`, *optional*):
                 RNG for reproducibility.
             num_inference_steps (`int`, defaults to `50`):
@@ -397,31 +250,34 @@ class JiTPipeline(DiffusionPipeline):
             return_dict (`bool`, *optional*, defaults to `True`):
                 Return [`ImagePipelineOutput`] if True.
         """
-        solver = sampling_method or self.scheduler.config.solver
-        if solver not in {"heun", "euler"}:
-            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
         if num_inference_steps < 2:
             raise ValueError("num_inference_steps must be >= 2.")
 
-        if t_eps is not None:
-            self.scheduler.register_to_config(t_eps=t_eps)
-
         class_label_ids = self._normalize_class_labels(class_labels)
         do_classifier_free_guidance = guidance_scale is not None and guidance_scale > 1.0
 
         batch_size = len(class_label_ids)
         image_size = int(self.transformer.config.sample_size)
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height <= 0 or width <= 0:
+            raise ValueError("height and width must be positive integers.")
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
         channels = int(self.transformer.config.in_channels)
         null_class_val = int(self.transformer.config.num_classes)
 
         if guidance_scale is None:
             guidance_scale = 1.0
         if noise_scale is None:
-            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(image_size, 1.0)
+            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)
 
         latents = (
             randn_tensor(
-                shape=(batch_size, channels, image_size, image_size),
+                shape=(batch_size, channels, height, width),
                 generator=generator,
                 device=self._execution_device,
                 dtype=self.transformer.dtype,
@@ -433,17 +289,47 @@ class JiTPipeline(DiffusionPipeline):
         class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
         class_null = torch.full_like(class_labels_t, null_class_val)
 
-        latents = self._run_sampler(
-            latents,
-            class_labels_t,
-            class_null,
-            num_inference_steps,
-            do_classifier_free_guidance,
-            guidance_scale,
-            guidance_interval_min,
-            guidance_interval_max,
-            solver,
-        )
+        if do_classifier_free_guidance:
+            class_labels_input = torch.cat([class_labels_t, class_null], dim=0)
+        else:
+            class_labels_input = class_labels_t
+
+        self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
+        for t in self.progress_bar(self.scheduler.timesteps):
+            step_index = self.scheduler.index_for_timestep(t, self.scheduler.timesteps)
+            sigma = self.scheduler.sigmas[step_index].to(device=latents.device, dtype=latents.dtype)
+            sigma = sigma.clamp_min(t_eps)
+            t_flow = (1.0 - sigma).clamp(0.0, 1.0)
+
+            if do_classifier_free_guidance:
+                latent_model_input = torch.cat([latents, latents], dim=0)
+            else:
+                latent_model_input = latents
+
+            timesteps = t_flow.flatten().expand(latent_model_input.shape[0])
+            x_pred = self.transformer(
+                latent_model_input,
+                timestep=timesteps,
+                class_labels=class_labels_input,
+                interpolate_pos_encoding=interpolate_pos_encoding,
+            ).sample
+
+            if do_classifier_free_guidance:
+                x_cond, x_uncond = x_pred.chunk(2, dim=0)
+                interval_mask = t_flow < guidance_interval_max
+                if guidance_interval_min != 0.0:
+                    interval_mask = interval_mask & (t_flow > guidance_interval_min)
+                scale = torch.where(
+                    interval_mask,
+                    torch.tensor(guidance_scale, device=latents.device, dtype=latents.dtype),
+                    torch.tensor(1.0, device=latents.device, dtype=latents.dtype),
+                )
+                x_pred = x_uncond + scale * (x_cond - x_uncond)
+
+            sigma = sigma.reshape(*([1] * (latents.ndim - 1)))
+            # JiT predicts x0; scheduler integrates in sigma space: dz/dsigma = -(x0 - z) / sigma.
+            model_output = -(x_pred - latents) / sigma
+            latents = self.scheduler.step(model_output, t, latents).prev_sample
 
         images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
         if output_type == "pt":

JiT-L-32/scheduler/scheduler_config.json

@@ -1,7 +1,6 @@
 {
-  "_class_name": "JiTScheduler",
+  "_class_name": "FlowMatchHeunDiscreteScheduler",
   "_diffusers_version": "0.36.0",
   "num_train_timesteps": 1000,
-  "t_eps": 0.05,
-  "solver": "heun"
+  "shift": 4.0
 }

JiT-L-32/transformer/jit_transformer_2d.py

@@ -68,38 +68,58 @@ class JiTRotaryEmbedding(nn.Module):
         num_cls_token=0,
     ):
         super().__init__()
-        if custom_freqs is not None:
-            freqs = custom_freqs
-        else:
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-
+        self.dim = dim
+        self.pt_seq_len = pt_seq_len
+        self.theta = theta
+        self.num_cls_token = num_cls_token
+        self.custom_freqs = custom_freqs
         if ft_seq_len is None:
             ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = freqs.repeat_interleave(2, dim=-1)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])  # [N_img, D]
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-
-            # prepend in-context cls token
-            _, D = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, D, dtype=cos_img.dtype)
-            sin_pad = torch.zeros(num_cls_token, D, dtype=sin_img.dtype)
-
-            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
-            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
+        self._cached_hw = None
+        cos, sin = self._build_freqs(ft_seq_len, ft_seq_len, device=torch.device("cpu"))
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+        self._cached_hw = (ft_seq_len, ft_seq_len)
+
+    def _build_freqs(self, height, width, device):
+        if self.custom_freqs is not None:
+            freqs = self.custom_freqs.to(device=device, dtype=torch.float32)
         else:
-            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
-            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
-
-    def forward(self, t):
+            freqs = 1.0 / (
+                self.theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)[: (self.dim // 2)] / self.dim)
+            )
+
+        t_h = torch.arange(height, device=device, dtype=torch.float32) / height * self.pt_seq_len
+        t_w = torch.arange(width, device=device, dtype=torch.float32) / width * self.pt_seq_len
+        freqs_h = torch.einsum("..., f -> ... f", t_h, freqs).repeat_interleave(2, dim=-1)
+        freqs_w = torch.einsum("..., f -> ... f", t_w, freqs).repeat_interleave(2, dim=-1)
+        freqs_2d = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
+        freqs_flat = freqs_2d.view(-1, freqs_2d.shape[-1])
+        cos_img = freqs_flat.cos()
+        sin_img = freqs_flat.sin()
+        if self.num_cls_token > 0:
+            _, dim_freq = cos_img.shape
+            cos_pad = torch.ones(self.num_cls_token, dim_freq, dtype=cos_img.dtype, device=device)
+            sin_pad = torch.zeros(self.num_cls_token, dim_freq, dtype=sin_img.dtype, device=device)
+            cos_img = torch.cat([cos_pad, cos_img], dim=0)
+            sin_img = torch.cat([sin_pad, sin_img], dim=0)
+        return cos_img, sin_img
+
+    def forward(self, t, height=None, width=None):
         # Applied on (batch, seq_len, heads, head_dim) tensors from attention.
         seq_len = t.shape[1]
+        if height is None or width is None:
+            image_tokens = seq_len - self.num_cls_token
+            size = int(image_tokens**0.5)
+            if size * size != image_tokens:
+                raise ValueError(
+                    f"Cannot infer square token grid from sequence length {seq_len} with {self.num_cls_token} class tokens."
+                )
+            height = size
+            width = size
+        if self._cached_hw != (height, width) or self.freqs_cos.device != t.device:
+            self.freqs_cos, self.freqs_sin = self._build_freqs(height, width, device=t.device)
+            self._cached_hw = (height, width)
         freqs_cos = self.freqs_cos[:seq_len].to(t.dtype)
         freqs_sin = self.freqs_sin[:seq_len].to(t.dtype)
 
@@ -195,7 +215,7 @@ class JiTAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
 
-    def forward(self, x, rope=None):
+    def forward(self, x, rope=None, grid_height=None, grid_width=None):
         B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
@@ -206,8 +226,8 @@ class JiTAttention(nn.Module):
         if rope is not None:
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
-            q = rope(q)
-            k = rope(k)
+            q = rope(q, height=grid_height, width=grid_width)
+            k = rope(k, height=grid_height, width=grid_width)
             q = q.transpose(1, 2)
             k = k.transpose(1, 2)
 
@@ -254,7 +274,7 @@ class JiTBlock(nn.Module):
         self.act = nn.SiLU()
         self.adaLN_modulation = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
 
-    def forward(self, x, c, feat_rope=None):
+    def forward(self, x, c, feat_rope=None, grid_height=None, grid_width=None):
         # Apply activation
         c = self.act(c)
 
@@ -263,7 +283,7 @@ class JiTBlock(nn.Module):
         # Attention block
         norm_x = self.norm1(x)
         modulated_x = modulate(norm_x, shift_msa, scale_msa)
-        attn_out = self.attn(modulated_x, rope=feat_rope)
+        attn_out = self.attn(modulated_x, rope=feat_rope, grid_height=grid_height, grid_width=grid_width)
         x = x + gate_msa.unsqueeze(1) * attn_out
 
         # MLP block
@@ -437,11 +457,30 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         self.act_final = nn.SiLU()
         self.adaLN_modulation_final = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
 
+    def _get_patch_grid(self, hidden_states):
+        height, width = hidden_states.shape[-2:]
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"Input size {(height, width)} must be divisible by patch_size={self.patch_size}."
+            )
+        return height // self.patch_size, width // self.patch_size
+
+    def _interpolate_pos_encoding(self, tokens, grid_height, grid_width):
+        num_tokens = grid_height * grid_width
+        if self.pos_embed.shape[1] == num_tokens:
+            return self.pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+        base_size = int(self.pos_embed.shape[1] ** 0.5)
+        pos_embed = self.pos_embed.reshape(1, base_size, base_size, self.hidden_size).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(grid_height, grid_width), mode="bicubic", align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(1, num_tokens, self.hidden_size)
+        return pos_embed.to(device=tokens.device, dtype=tokens.dtype)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         class_labels: torch.LongTensor,
+        interpolate_pos_encoding: bool = True,
         return_dict: bool = True,
     ):
 
@@ -454,8 +493,19 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         c = t_emb + y_emb
 
         # Patch Embed
+        grid_height, grid_width = self._get_patch_grid(hidden_states)
         x = self.x_embedder(hidden_states)
-        x = x + self.pos_embed.to(x.dtype)
+        if interpolate_pos_encoding:
+            pos_embed = self._interpolate_pos_encoding(x, grid_height, grid_width)
+        else:
+            expected_tokens = grid_height * grid_width
+            if self.pos_embed.shape[1] != expected_tokens:
+                raise ValueError(
+                    f"pos_embed token count {self.pos_embed.shape[1]} does not match input token count {expected_tokens}. "
+                    "Enable interpolate_pos_encoding for dynamic resolutions."
+                )
+            pos_embed = self.pos_embed.to(device=x.device, dtype=x.dtype)
+        x = x + pos_embed
 
         # Blocks
         for i, block in enumerate(self.blocks):
@@ -467,15 +517,23 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
             rope = self.feat_rope if i < self.in_context_start else self.feat_rope_incontext
 
             if self.training and self.gradient_checkpointing:
+                def custom_forward(current_x, current_c):
+                    return block(
+                        current_x,
+                        current_c,
+                        feat_rope=rope,
+                        grid_height=grid_height,
+                        grid_width=grid_width,
+                    )
+
                 x = torch.utils.checkpoint.checkpoint(
-                    block,
+                    custom_forward,
                     x,
                     c,
-                    rope,
                     use_reentrant=False,
                 )
             else:
-                x = block(x, c, feat_rope=rope)
+                x = block(x, c, feat_rope=rope, grid_height=grid_height, grid_width=grid_width)
 
         # Slice off in-context tokens
         if self.in_context_len > 0:
@@ -489,10 +547,11 @@ class JiTTransformer2DModel(ModelMixin, ConfigMixin):
         x = self.linear_final(x)
 
         # Unpatchify
-        h = w = int(x.shape[1] ** 0.5)
-        x = x.reshape(shape=(x.shape[0], h, w, self.patch_size, self.patch_size, self.out_channels))
+        x = x.reshape(shape=(x.shape[0], grid_height, grid_width, self.patch_size, self.patch_size, self.out_channels))
         x = torch.einsum("nhwpqc->nchpwq", x)
-        output = x.reshape(shape=(x.shape[0], self.out_channels, h * self.patch_size, w * self.patch_size))
+        output = x.reshape(
+            shape=(x.shape[0], self.out_channels, grid_height * self.patch_size, grid_width * self.patch_size)
+        )
 
         if not return_dict:
             return (output,)

README.md

@@ -19,17 +19,17 @@ language:
 Native diffusers implementation of **JiT** (Just image Transformer). Each variant folder is self-contained:
 
 - `pipeline.py` — `JiTPipeline`
-- `scheduler/scheduling_jit.py` — `JiTScheduler` (linear `t in [0, 1]`, Heun/Euler)
+- `scheduler/scheduler_config.json` — `FlowMatchHeunDiscreteScheduler` config (default `shift=4.0`)
 - `transformer/jit_transformer_2d.py` — `JiTTransformer2DModel`
 
-Shared ImageNet-1k labels live in [`labels/`](labels/) at the repo root (not duplicated per variant).
+The pipeline now supports dynamic inference resolution in `__call__` with positional interpolation.
 
 No separate `jit_diffusers` package; only PyPI `diffusers` plus local custom code in the variant directory.
 
 ## Available checkpoints
 
 | Checkpoint | Path | Resolution | Recommended CFG |
-|---|---|---|---|
+| --- | --- | --- | --- |
 | JiT-B/16 | `./JiT-B-16` | 256×256 | 3.0 |
 | JiT-L/16 | `./JiT-L-16` | 256×256 | 2.4 |
 | JiT-H/16 | `./JiT-H-16` | 256×256 | 2.2 |
@@ -39,42 +39,52 @@ No separate `jit_diffusers` package; only PyPI `diffusers` plus local custom cod
 
 ## ImageNet class labels
 
-| File | Direction | Format |
-|---|---|---|
-| `labels/id2label_en.json` | id → English | comma-separated synonyms, e.g. `"207": "golden retriever"` |
-| `labels/id2label_cn.json` | id → Chinese | comma-separated synonyms, e.g. `"207": "金毛猎犬"` |
+Each variant keeps an English `id2label` map directly in its own `model_index.json` (DiT-style).
 
-- `pipe.id2label` / `pipe.id2label_cn` — inspect id → label correspondence
-- `pipe.labels` / `pipe.labels_cn` — reverse maps (synonym → id), sorted for browsing
-- `pipe.get_label_ids("golden retriever")` or `pipe.get_label_ids("金毛猎犬", lang="cn")`
+- `pipe.id2label` — inspect id → English label correspondence
+- `pipe.labels` — reverse map (English synonym → id), sorted for browsing
+- `pipe.get_label_ids("golden retriever")`
 - `pipe(class_labels="golden retriever", ...)` — string labels resolved automatically
 
+Chinese labels are preserved in the main source repo under `src/labels/id2label_cn.json` for reference.
+
 ## Inference
 
+Run the bundled demo script from the repo root:
+
+```bash
+python demo_inference.py
+```
+
+This writes `demo.png` using `JiT-H-32` with the settings below.
+
 ```python
-from diffusers import DiffusionPipeline
+from pathlib import Path
+from diffusers import DiffusionPipeline, FlowMatchHeunDiscreteScheduler
 import torch
 
+model_dir = Path("./JiT-H-32")
 pipe = DiffusionPipeline.from_pretrained(
-    "./JiT-H-32",
-    trust_remote_code=True,
+    str(model_dir),
+    custom_pipeline=str(model_dir / "pipeline.py"),
 )
+pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(pipe.scheduler.config, shift=4.0)
 pipe.to("cuda")
-pipe.transformer.to(dtype=torch.bfloat16)
 
 # Numeric or human-readable labels
 print(pipe.id2label[207])
 print(pipe.get_label_ids("golden retriever"))
 
 generator = torch.Generator(device="cuda").manual_seed(42)
-images = pipe(
+image = pipe(
     class_labels="golden retriever",
     num_inference_steps=50,
     guidance_scale=2.3,
-    sampling_method="heun",
     generator=generator,
-).images
-images[0].save("output.png")
+).images[0]
+image.save("demo.png")
 ```
 
+`height` and `width` default to the checkpoint's native resolution when omitted.
+
 Load a **variant subfolder** (e.g. `./JiT-H-32`), not the repo root.

demo_inference.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""Generate a demo image with JiT-H-32."""
+
+from pathlib import Path
+
+import torch
+from diffusers import DiffusionPipeline, FlowMatchHeunDiscreteScheduler
+
+REPO_ROOT = Path(__file__).resolve().parent
+MODEL_DIR = REPO_ROOT / "JiT-H-32"
+OUTPUT_PATH = REPO_ROOT / "demo.png"
+
+
+def main() -> None:
+    pipe = DiffusionPipeline.from_pretrained(
+        str(MODEL_DIR),
+        custom_pipeline=str(MODEL_DIR / "pipeline.py"),
+        torch_dtype=torch.bfloat16,
+    )
+    pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(pipe.scheduler.config, shift=4.0)
+    pipe.to("cuda")
+
+    print(pipe.id2label[207])
+    print(pipe.get_label_ids("golden retriever"))
+
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    image = pipe(
+        class_labels="golden retriever",
+        num_inference_steps=50,
+        guidance_scale=2.3,
+        generator=generator,
+    ).images[0]
+    image.save(OUTPUT_PATH)
+    print(f"Saved demo image to {OUTPUT_PATH}")
+
+
+if __name__ == "__main__":
+    main()