Upload folder using huggingface_hub

var/D3HR/DiT-XL-2-256x256.pt

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

var/D3HR/DiT-XL/.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

var/D3HR/DiT-XL/README.md

@@ -0,0 +1,10 @@
+---
+license: cc-by-nc-4.0
+---
+
+# Scalable Diffusion Models with Transformers (DiT)
+
+## Abstract
+
+We train latent diffusion models, replacing the commonly-used U-Net backbone with a transformer that operates on latent patches. We analyze the scalability of our Diffusion Transformers (DiTs) through the lens of forward pass complexity as measured by Gflops. We find that DiTs with higher Gflops---through increased transformer depth/width or increased number of input tokens---consistently have lower FID. In addition to good scalability properties, our DiT-XL/2 models outperform all prior diffusion models on the class-conditional ImageNet 512×512 and 256×256 benchmarks, achieving a state-of-the-art FID of 2.27 on the latter.
+

var/D3HR/DiT-XL/model_index.json

@@ -0,0 +1,1018 @@
+{
+  "_class_name": "DiTPipeline",
+  "_diffusers_version": "0.12.0.dev0",
+  "scheduler": [
+    "diffusers",
+    "DDIMScheduler"
+  ],
+  "transformer": [
+    "diffusers",
+    "Transformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ],
+  "id2label": {
+    "0": "tench, Tinca tinca",
+    "1": "goldfish, Carassius auratus",
+    "2": "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
+    "3": "tiger shark, Galeocerdo cuvieri",
+    "4": "hammerhead, hammerhead shark",
+    "5": "electric ray, crampfish, numbfish, torpedo",
+    "6": "stingray",
+    "7": "cock",
+    "8": "hen",
+    "9": "ostrich, Struthio camelus",
+    "10": "brambling, Fringilla montifringilla",
+    "11": "goldfinch, Carduelis carduelis",
+    "12": "house finch, linnet, Carpodacus mexicanus",
+    "13": "junco, snowbird",
+    "14": "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+    "15": "robin, American robin, Turdus migratorius",
+    "16": "bulbul",
+    "17": "jay",
+    "18": "magpie",
+    "19": "chickadee",
+    "20": "water ouzel, dipper",
+    "21": "kite",
+    "22": "bald eagle, American eagle, Haliaeetus leucocephalus",
+    "23": "vulture",
+    "24": "great grey owl, great gray owl, Strix nebulosa",
+    "25": "European fire salamander, Salamandra salamandra",
+    "26": "common newt, Triturus vulgaris",
+    "27": "eft",
+    "28": "spotted salamander, Ambystoma maculatum",
+    "29": "axolotl, mud puppy, Ambystoma mexicanum",
+    "30": "bullfrog, Rana catesbeiana",
+    "31": "tree frog, tree-frog",
+    "32": "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+    "33": "loggerhead, loggerhead turtle, Caretta caretta",
+    "34": "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
+    "35": "mud turtle",
+    "36": "terrapin",
+    "37": "box turtle, box tortoise",
+    "38": "banded gecko",
+    "39": "common iguana, iguana, Iguana iguana",
+    "40": "American chameleon, anole, Anolis carolinensis",
+    "41": "whiptail, whiptail lizard",
+    "42": "agama",
+    "43": "frilled lizard, Chlamydosaurus kingi",
+    "44": "alligator lizard",
+    "45": "Gila monster, Heloderma suspectum",
+    "46": "green lizard, Lacerta viridis",
+    "47": "African chameleon, Chamaeleo chamaeleon",
+    "48": "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
+    "49": "African crocodile, Nile crocodile, Crocodylus niloticus",
+    "50": "American alligator, Alligator mississipiensis",
+    "51": "triceratops",
+    "52": "thunder snake, worm snake, Carphophis amoenus",
+    "53": "ringneck snake, ring-necked snake, ring snake",
+    "54": "hognose snake, puff adder, sand viper",
+    "55": "green snake, grass snake",
+    "56": "king snake, kingsnake",
+    "57": "garter snake, grass snake",
+    "58": "water snake",
+    "59": "vine snake",
+    "60": "night snake, Hypsiglena torquata",
+    "61": "boa constrictor, Constrictor constrictor",
+    "62": "rock python, rock snake, Python sebae",
+    "63": "Indian cobra, Naja naja",
+    "64": "green mamba",
+    "65": "sea snake",
+    "66": "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+    "67": "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+    "68": "sidewinder, horned rattlesnake, Crotalus cerastes",
+    "69": "trilobite",
+    "70": "harvestman, daddy longlegs, Phalangium opilio",
+    "71": "scorpion",
+    "72": "black and gold garden spider, Argiope aurantia",
+    "73": "barn spider, Araneus cavaticus",
+    "74": "garden spider, Aranea diademata",
+    "75": "black widow, Latrodectus mactans",
+    "76": "tarantula",
+    "77": "wolf spider, hunting spider",
+    "78": "tick",
+    "79": "centipede",
+    "80": "black grouse",
+    "81": "ptarmigan",
+    "82": "ruffed grouse, partridge, Bonasa umbellus",
+    "83": "prairie chicken, prairie grouse, prairie fowl",
+    "84": "peacock",
+    "85": "quail",
+    "86": "partridge",
+    "87": "African grey, African gray, Psittacus erithacus",
+    "88": "macaw",
+    "89": "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+    "90": "lorikeet",
+    "91": "coucal",
+    "92": "bee eater",
+    "93": "hornbill",
+    "94": "hummingbird",
+    "95": "jacamar",
+    "96": "toucan",
+    "97": "drake",
+    "98": "red-breasted merganser, Mergus serrator",
+    "99": "goose",
+    "100": "black swan, Cygnus atratus",
+    "101": "tusker",
+    "102": "echidna, spiny anteater, anteater",
+    "103": "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",
+    "104": "wallaby, brush kangaroo",
+    "105": "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",
+    "106": "wombat",
+    "107": "jellyfish",
+    "108": "sea anemone, anemone",
+    "109": "brain coral",
+    "110": "flatworm, platyhelminth",
+    "111": "nematode, nematode worm, roundworm",
+    "112": "conch",
+    "113": "snail",
+    "114": "slug",
+    "115": "sea slug, nudibranch",
+    "116": "chiton, coat-of-mail shell, sea cradle, polyplacophore",
+    "117": "chambered nautilus, pearly nautilus, nautilus",
+    "118": "Dungeness crab, Cancer magister",
+    "119": "rock crab, Cancer irroratus",
+    "120": "fiddler crab",
+    "121": "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica",
+    "122": "American lobster, Northern lobster, Maine lobster, Homarus americanus",
+    "123": "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish",
+    "124": "crayfish, crawfish, crawdad, crawdaddy",
+    "125": "hermit crab",
+    "126": "isopod",
+    "127": "white stork, Ciconia ciconia",
+    "128": "black stork, Ciconia nigra",
+    "129": "spoonbill",
+    "130": "flamingo",
+    "131": "little blue heron, Egretta caerulea",
+    "132": "American egret, great white heron, Egretta albus",
+    "133": "bittern",
+    "134": "crane",
+    "135": "limpkin, Aramus pictus",
+    "136": "European gallinule, Porphyrio porphyrio",
+    "137": "American coot, marsh hen, mud hen, water hen, Fulica americana",
+    "138": "bustard",
+    "139": "ruddy turnstone, Arenaria interpres",
+    "140": "red-backed sandpiper, dunlin, Erolia alpina",
+    "141": "redshank, Tringa totanus",
+    "142": "dowitcher",
+    "143": "oystercatcher, oyster catcher",
+    "144": "pelican",
+    "145": "king penguin, Aptenodytes patagonica",
+    "146": "albatross, mollymawk",
+    "147": "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus",
+    "148": "killer whale, killer, orca, grampus, sea wolf, Orcinus orca",
+    "149": "dugong, Dugong dugon",
+    "150": "sea lion",
+    "151": "Chihuahua",
+    "152": "Japanese spaniel",
+    "153": "Maltese dog, Maltese terrier, Maltese",
+    "154": "Pekinese, Pekingese, Peke",
+    "155": "Shih-Tzu",
+    "156": "Blenheim spaniel",
+    "157": "papillon",
+    "158": "toy terrier",
+    "159": "Rhodesian ridgeback",
+    "160": "Afghan hound, Afghan",
+    "161": "basset, basset hound",
+    "162": "beagle",
+    "163": "bloodhound, sleuthhound",
+    "164": "bluetick",
+    "165": "black-and-tan coonhound",
+    "166": "Walker hound, Walker foxhound",
+    "167": "English foxhound",
+    "168": "redbone",
+    "169": "borzoi, Russian wolfhound",
+    "170": "Irish wolfhound",
+    "171": "Italian greyhound",
+    "172": "whippet",
+    "173": "Ibizan hound, Ibizan Podenco",
+    "174": "Norwegian elkhound, elkhound",
+    "175": "otterhound, otter hound",
+    "176": "Saluki, gazelle hound",
+    "177": "Scottish deerhound, deerhound",
+    "178": "Weimaraner",
+    "179": "Staffordshire bullterrier, Staffordshire bull terrier",
+    "180": "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",
+    "181": "Bedlington terrier",
+    "182": "Border terrier",
+    "183": "Kerry blue terrier",
+    "184": "Irish terrier",
+    "185": "Norfolk terrier",
+    "186": "Norwich terrier",
+    "187": "Yorkshire terrier",
+    "188": "wire-haired fox terrier",
+    "189": "Lakeland terrier",
+    "190": "Sealyham terrier, Sealyham",
+    "191": "Airedale, Airedale terrier",
+    "192": "cairn, cairn terrier",
+    "193": "Australian terrier",
+    "194": "Dandie Dinmont, Dandie Dinmont terrier",
+    "195": "Boston bull, Boston terrier",
+    "196": "miniature schnauzer",
+    "197": "giant schnauzer",
+    "198": "standard schnauzer",
+    "199": "Scotch terrier, Scottish terrier, Scottie",
+    "200": "Tibetan terrier, chrysanthemum dog",
+    "201": "silky terrier, Sydney silky",
+    "202": "soft-coated wheaten terrier",
+    "203": "West Highland white terrier",
+    "204": "Lhasa, Lhasa apso",
+    "205": "flat-coated retriever",
+    "206": "curly-coated retriever",
+    "207": "golden retriever",
+    "208": "Labrador retriever",
+    "209": "Chesapeake Bay retriever",
+    "210": "German short-haired pointer",
+    "211": "vizsla, Hungarian pointer",
+    "212": "English setter",
+    "213": "Irish setter, red setter",
+    "214": "Gordon setter",
+    "215": "Brittany spaniel",
+    "216": "clumber, clumber spaniel",
+    "217": "English springer, English springer spaniel",
+    "218": "Welsh springer spaniel",
+    "219": "cocker spaniel, English cocker spaniel, cocker",
+    "220": "Sussex spaniel",
+    "221": "Irish water spaniel",
+    "222": "kuvasz",
+    "223": "schipperke",
+    "224": "groenendael",
+    "225": "malinois",
+    "226": "briard",
+    "227": "kelpie",
+    "228": "komondor",
+    "229": "Old English sheepdog, bobtail",
+    "230": "Shetland sheepdog, Shetland sheep dog, Shetland",
+    "231": "collie",
+    "232": "Border collie",
+    "233": "Bouvier des Flandres, Bouviers des Flandres",
+    "234": "Rottweiler",
+    "235": "German shepherd, German shepherd dog, German police dog, alsatian",
+    "236": "Doberman, Doberman pinscher",
+    "237": "miniature pinscher",
+    "238": "Greater Swiss Mountain dog",
+    "239": "Bernese mountain dog",
+    "240": "Appenzeller",
+    "241": "EntleBucher",
+    "242": "boxer",
+    "243": "bull mastiff",
+    "244": "Tibetan mastiff",
+    "245": "French bulldog",
+    "246": "Great Dane",
+    "247": "Saint Bernard, St Bernard",
+    "248": "Eskimo dog, husky",
+    "249": "malamute, malemute, Alaskan malamute",
+    "250": "Siberian husky",
+    "251": "dalmatian, coach dog, carriage dog",
+    "252": "affenpinscher, monkey pinscher, monkey dog",
+    "253": "basenji",
+    "254": "pug, pug-dog",
+    "255": "Leonberg",
+    "256": "Newfoundland, Newfoundland dog",
+    "257": "Great Pyrenees",
+    "258": "Samoyed, Samoyede",
+    "259": "Pomeranian",
+    "260": "chow, chow chow",
+    "261": "keeshond",
+    "262": "Brabancon griffon",
+    "263": "Pembroke, Pembroke Welsh corgi",
+    "264": "Cardigan, Cardigan Welsh corgi",
+    "265": "toy poodle",
+    "266": "miniature poodle",
+    "267": "standard poodle",
+    "268": "Mexican hairless",
+    "269": "timber wolf, grey wolf, gray wolf, Canis lupus",
+    "270": "white wolf, Arctic wolf, Canis lupus tundrarum",
+    "271": "red wolf, maned wolf, Canis rufus, Canis niger",
+    "272": "coyote, prairie wolf, brush wolf, Canis latrans",
+    "273": "dingo, warrigal, warragal, Canis dingo",
+    "274": "dhole, Cuon alpinus",
+    "275": "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
+    "276": "hyena, hyaena",
+    "277": "red fox, Vulpes vulpes",
+    "278": "kit fox, Vulpes macrotis",
+    "279": "Arctic fox, white fox, Alopex lagopus",
+    "280": "grey fox, gray fox, Urocyon cinereoargenteus",
+    "281": "tabby, tabby cat",
+    "282": "tiger cat",
+    "283": "Persian cat",
+    "284": "Siamese cat, Siamese",
+    "285": "Egyptian cat",
+    "286": "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor",
+    "287": "lynx, catamount",
+    "288": "leopard, Panthera pardus",
+    "289": "snow leopard, ounce, Panthera uncia",
+    "290": "jaguar, panther, Panthera onca, Felis onca",
+    "291": "lion, king of beasts, Panthera leo",
+    "292": "tiger, Panthera tigris",
+    "293": "cheetah, chetah, Acinonyx jubatus",
+    "294": "brown bear, bruin, Ursus arctos",
+    "295": "American black bear, black bear, Ursus americanus, Euarctos americanus",
+    "296": "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus",
+    "297": "sloth bear, Melursus ursinus, Ursus ursinus",
+    "298": "mongoose",
+    "299": "meerkat, mierkat",
+    "300": "tiger beetle",
+    "301": "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle",
+    "302": "ground beetle, carabid beetle",
+    "303": "long-horned beetle, longicorn, longicorn beetle",
+    "304": "leaf beetle, chrysomelid",
+    "305": "dung beetle",
+    "306": "rhinoceros beetle",
+    "307": "weevil",
+    "308": "fly",
+    "309": "bee",
+    "310": "ant, emmet, pismire",
+    "311": "grasshopper, hopper",
+    "312": "cricket",
+    "313": "walking stick, walkingstick, stick insect",
+    "314": "cockroach, roach",
+    "315": "mantis, mantid",
+    "316": "cicada, cicala",
+    "317": "leafhopper",
+    "318": "lacewing, lacewing fly",
+    "319": "dragonfly, darning needle, devil's 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's 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": "carpenter's 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, jeweler's 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, pj's, 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, carpenter's plane, woodworking plane",
+    "727": "planetarium",
+    "728": "plastic bag",
+    "729": "plate rack",
+    "730": "plow, plough",
+    "731": "plunger, plumber's 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": "potter's 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, spider's 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 lady's 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"
+  }
+}

var/D3HR/DiT-XL/scheduler/scheduler_config.json

@@ -0,0 +1,13 @@
+{
+  "_class_name": "DDIMScheduler",
+  "_diffusers_version": "0.12.0.dev0",
+  "beta_end": 0.02,
+  "beta_schedule": "linear",
+  "beta_start": 0.0001,
+  "clip_sample": false,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "set_alpha_to_one": true,
+  "steps_offset": 0,
+  "trained_betas": null
+}

var/D3HR/DiT-XL/transformer/config.json

@@ -0,0 +1,23 @@
+{
+  "_class_name": "Transformer2DModel",
+  "_diffusers_version": "0.12.0.dev0",
+  "activation_fn": "gelu-approximate",
+  "attention_bias": true,
+  "attention_head_dim": 72,
+  "cross_attention_dim": null,
+  "dropout": 0.0,
+  "in_channels": 4,
+  "norm_elementwise_affine": false,
+  "norm_num_groups": 32,
+  "norm_type": "ada_norm_zero",
+  "num_attention_heads": 16,
+  "num_embeds_ada_norm": 1000,
+  "num_layers": 28,
+  "num_vector_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 8,
+  "patch_size": 2,
+  "sample_size": 32,
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

var/D3HR/DiT-XL/transformer/diffusion_pytorch_model.bin

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

var/D3HR/DiT-XL/vae/config.json

@@ -0,0 +1,30 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.12.0.dev0",
+  "_name_or_path": "stabilityai/sd-vae-ft-ema",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

var/D3HR/DiT-XL/vae/diffusion_pytorch_model.bin

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

var/D3HR/README.md

@@ -0,0 +1,98 @@
+# Taming Diffusion for Dataset Distillation with High Representativeness (ICML 2025)
+
+This repository is the official implementation of the paper:
+
+[**Taming Diffusion for Dataset Distillation with High Representativeness**](https://www.arxiv.org/pdf/2505.18399)
+[*Lin Zhao*](https://lin-zhao-resolve.github.io/),
+[*Yushu Wu*](https://wuyushuwys.github.io/),
+[*Xinru Jiang*](https://oshikaka.github.io/),
+[*Jianyang Gu*](https://vimar-gu.github.io/),
+[*Yanzhi Wang*](https://coe.northeastern.edu/people/wang-yanzhi/),
+[*Xiaolin Xu*](https://www.xiaolinxu.com/),
+[*Pu Zhao*](https://puzhao.info/),
+[*Xue Lin*](https://coe.northeastern.edu/people/lin-xue/),
+ICML, 2025.
+
+<div align=center>
+<img width=85% src="./imgs/framework.jpg"/>
+</div>
+
+## Usage
+
+1. [Distilled Datasets](#distilled-datasets)
+2. [Setup](#setup)
+3. [Step1: DDIM inversion and distribution matching](#step1-ddim-inversion-and-distribution-matching)
+4. [Step2: Group sampling](#step2-group-sampling)
+6. [Evaluation](#evaluation)
+
+
+
+## Distilled Datasets
+We provide distilled datasets with different IPCs generated by our method on Huggingface🤗! [*Imagenet-1K*](https://www.image-net.org/), [*Tiny-Imagenet*](https://www.kaggle.com/c/tiny-imagenet), [*CIFAR10*](https://www.cs.toronto.edu/~kriz/cifar.html), [*CIFAR100*](https://www.cs.toronto.edu/~kriz/cifar.html) datasets for users to use directly. 
+
+🔥Distilled datasets for Imagenet-1K: [10IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/imagenet1k_10ipc), [50IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/imagenet1k_50ipc)
+
+🔥Distilled datasets for Tiny-Imagenet: [10IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/tinyimagenet_10ipc), [50IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/tinyimagenet_50ipc)
+
+🔥Distilled datasets for CIFAR10: [10IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/cifar10_10ipc), [50IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/cifar10_50ipc)
+
+🔥Distilled datasets for CIFAR100: [10IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/cifar100_10ipc), [50IPC](https://huggingface.co/datasets/lin-zhao-resoLve/D3HR/tree/main/cifar100_50ipc)
+
+Besides, if you want to use the D3HR to generate distilled datasets by yourself, run the following steps:
+
+## Setup
+
+To install the required dependencies, use the following commands:
+
+```bash
+conda create -n D3HR python=3.10
+conda activate D3HR
+cd D3HR
+pip install -e .
+```
+
+## Step1: DDIM inversion and distribution matching
+
+### (1) load pretrained model
+For Imagenet-1K dataset, you can just use the pretrained [DiT](https://github.com/facebookresearch/DiT) model in huggingface:
+```bash
+huggingface-cli download facebook/DiT-XL-2-256 --local-dir <your_local_path>
+```
+For other datasets, you must first fine-tune the pretrained DiT model on the dataset ([github repo](https://github.com/facebookresearch/DiT)), then continue. 
+
+### (2) perform DDIM inversion and distribution matching to obtain the statistic information
+```bash
+sh generation/dit_inversion_save_statistic.sh
+```
+Note: By default, we store the results at 15 timesteps (23 < t < 39) to support the experiments in Section 6.2.
+
+## Step2: Group sampling
+```bash
+sh generation/group_sampling.sh
+```
+
+
+## Evaluation
+```bash
+sh validation/validate.sh
+```
+Note: The .sh script includes several configuration options—select the one that best fits your needs.
+
+## Acknowledgement
+This project is mainly developed based on:
+[DiT](https://github.com/facebookresearch/DiT)
+
+
+## Contact
+If you have any questions, please contact zhao.lin1@northeastern.edu.
+
+## Citation
+If you find our work useful, please cite:
+
+```BiBTeX
+@inproceedings{zhaotaming,
+  title={Taming Diffusion for Dataset Distillation with High Representativeness},
+  author={Zhao, Lin and Wu, Yushu and Jiang, Xinru and Gu, Jianyang and Wang, Yanzhi and Xu, Xiaolin and Zhao, Pu and Lin, Xue},
+  booktitle={Forty-second International Conference on Machine Learning}
+}
+```

var/D3HR/ds_inf/imagenet1k_train.txt

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

var/D3HR/ds_inf/tiny-imagenet-mapping.txt

@@ -0,0 +1,200 @@
+n02124075	285
+n04067472	758
+n04540053	890
+n04099969	765
+n07749582	951
+n01641577	30
+n02802426	430
+n09246464	972
+n07920052	967
+n03970156	731
+n03891332	704
+n02106662	235
+n03201208	532
+n02279972	323
+n02132136	294
+n04146614	779
+n07873807	963
+n02364673	338
+n04507155	879
+n03854065	687
+n03838899	683
+n03733131	645
+n01443537	1
+n07875152	964
+n03544143	604
+n09428293	978
+n03085013	508
+n02437312	354
+n07614500	928
+n03804744	677
+n04265275	811
+n02963159	474
+n02486410	372
+n01944390	113
+n09256479	973
+n02058221	146
+n04275548	815
+n02321529	329
+n02769748	414
+n02099712	208
+n07695742	932
+n02056570	145
+n02281406	325
+n01774750	76
+n02509815	387
+n03983396	737
+n07753592	954
+n04254777	806
+n02233338	314
+n04008634	744
+n02823428	440
+n02236044	315
+n03393912	565
+n07583066	924
+n04074963	761
+n01629819	25
+n09332890	975
+n02481823	367
+n03902125	707
+n03404251	568
+n09193705	970
+n03637318	619
+n04456115	862
+n02666196	398
+n03796401	675
+n02795169	427
+n02123045	281
+n01855672	99
+n01882714	105
+n02917067	466
+n02988304	485
+n04398044	849
+n02843684	448
+n02423022	353
+n02669723	400
+n04465501	866
+n02165456	301
+n03770439	655
+n02099601	207
+n04486054	873
+n02950826	471
+n03814639	678
+n04259630	808
+n03424325	570
+n02948072	470
+n03179701	526
+n03400231	567
+n02206856	309
+n03160309	525
+n01984695	123
+n03977966	734
+n03584254	605
+n04023962	747
+n02814860	437
+n01910747	107
+n04596742	909
+n03992509	739
+n04133789	774
+n03937543	720
+n02927161	467
+n01945685	114
+n02395406	341
+n02125311	286
+n03126707	517
+n04532106	887
+n02268443	319
+n02977058	480
+n07734744	947
+n03599486	612
+n04562935	900
+n03014705	492
+n04251144	801
+n04356056	837
+n02190166	308
+n03670208	627
+n02002724	128
+n02074367	149
+n04285008	817
+n04560804	899
+n04366367	839
+n02403003	345
+n07615774	929
+n04501370	877
+n03026506	496
+n02906734	462
+n01770393	71
+n04597913	910
+n03930313	716
+n04118538	768
+n04179913	786
+n04311004	821
+n02123394	283
+n04070727	760
+n02793495	425
+n02730930	411
+n02094433	187
+n04371430	842
+n04328186	826
+n03649909	621
+n04417672	853
+n03388043	562
+n01774384	75
+n02837789	445
+n07579787	923
+n04399382	850
+n02791270	424
+n03089624	509
+n02814533	436
+n04149813	781
+n07747607	950
+n03355925	557
+n01983481	122
+n04487081	874
+n03250847	542
+n03255030	543
+n02892201	458
+n02883205	457
+n03100240	511
+n02415577	349
+n02480495	365
+n01698640	50
+n01784675	79
+n04376876	845
+n03444034	573
+n01917289	109
+n01950731	115
+n03042490	500
+n07711569	935
+n04532670	888
+n03763968	652
+n07768694	957
+n02999410	488
+n03617480	614
+n06596364	917
+n01768244	69
+n02410509	347
+n03976657	733
+n01742172	61
+n03980874	735
+n02808440	435
+n02226429	311
+n02231487	313
+n02085620	151
+n01644900	32
+n02129165	291
+n02699494	406
+n03837869	682
+n02815834	438
+n07720875	945
+n02788148	421
+n02909870	463
+n03706229	635
+n07871810	962
+n03447447	576
+n02113799	267
+n12267677	988
+n03662601	625
+n02841315	447
+n07715103	938
+n02504458	386

var/D3HR/generation/dit_inversion_save_statistic.py

@@ -0,0 +1,437 @@
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from io import BytesIO
+from tqdm.auto import tqdm
+from matplotlib import pyplot as plt
+from torchvision import transforms as tfms
+from diffusers import StableDiffusionPipeline, DDIMScheduler, DiTPipeline
+import argparse
+import os
+from scipy import io
+from diffusers import DiTPipeline
+from diffusers.utils.torch_utils import randn_tensor
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import ipdb
+from torch.utils.data import Dataset
+import torchvision.transforms as transforms
+from itertools import islice
+import json
+
+# sample
+@torch.no_grad()
+def sample(
+    pipe,
+    class_labels,
+    start_step=0,
+    start_latents=None,
+    guidance_scale=4.0,
+    num_inference_steps=30,
+    do_classifier_free_guidance=True,
+    device=None,
+):
+    
+    batch_size = len(class_labels)
+    latent_size = pipe.transformer.config.sample_size
+    latent_channels = pipe.transformer.config.in_channels
+    if start_latents == None:
+        latents = randn_tensor(
+            shape=(batch_size, latent_channels, latent_size, latent_size),
+            generator=generator,
+            device=pipe._execution_device,
+            dtype=pipe.transformer.dtype,
+        )
+    else:
+        latents = start_latents.clone()
+    
+    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+
+    class_labels = torch.tensor(class_labels, device=device).reshape(-1)
+    class_null = torch.tensor([1000] * batch_size, device=device)
+    class_labels_input = torch.cat([class_labels, class_null], 0) if guidance_scale > 1 else class_labels
+    class_labels_input = class_labels_input.to(device)
+
+    # set step values
+    pipe.scheduler.set_timesteps(num_inference_steps)
+    
+    
+    for i in tqdm(range(start_step, num_inference_steps)):
+    
+        t = pipe.scheduler.timesteps[i]
+        
+        if do_classifier_free_guidance:
+            half = latent_model_input[: len(latent_model_input) // 2]
+            latent_model_input = torch.cat([half, half], dim=0)
+            
+        latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
+
+        timesteps = t
+        if not torch.is_tensor(timesteps):
+            is_mps = latent_model_input.device.type == "mps"
+            if isinstance(timesteps, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=latent_model_input.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(latent_model_input.device)
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(latent_model_input.shape[0])
+        latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
+        # predict noise model_output
+        noise_pred = pipe.transformer(
+            latent_model_input, timestep=timesteps, class_labels=class_labels_input
+        ).sample
+
+         # Perform guidance
+        if do_classifier_free_guidance:
+            # perform guidance
+            # noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            # noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        
+            eps, rest = noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
+            cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+
+            half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
+            eps = torch.cat([half_eps, half_eps], dim=0)
+
+            noise_pred = torch.cat([eps, rest], dim=1)
+
+        # learned sigma
+        if pipe.transformer.config.out_channels // 2 == latent_channels:
+        
+            model_output, _ = torch.split(noise_pred, latent_channels, dim=1)
+        else:
+            model_output = noise_pred
+
+        # calculate ddim:
+        prev_t = max(1, t.item() - (1000 // num_inference_steps))  # t-1
+        alpha_t = pipe.scheduler.alphas_cumprod[t.item()]
+        alpha_t_prev = pipe.scheduler.alphas_cumprod[prev_t]
+        predicted_x0 = (latent_model_input - (1 - alpha_t).sqrt() * model_output) / alpha_t.sqrt()
+        direction_pointing_to_xt = (1 - alpha_t_prev).sqrt() * model_output
+        latent_model_input = alpha_t_prev.sqrt() * predicted_x0 + direction_pointing_to_xt
+        # latent_model_input = pipe.scheduler.step(model_output, t, latent_model_input).prev_sample
+
+    if guidance_scale > 1:
+        latents, _ = latent_model_input.chunk(2, dim=0)
+    else:
+        latents = latent_model_input
+            
+    latents = 1 / pipe.vae.config.scaling_factor * latents
+    samples = pipe.vae.decode(latents).sample
+
+    samples = (samples / 2 + 0.5).clamp(0, 1)
+
+    # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+    samples = samples.cpu().permute(0, 2, 3, 1).float().numpy()
+
+    return samples
+
+    
+
+
+
+
+## Inversion
+@torch.no_grad()
+def invert(
+    pipe,
+    start_latents,
+    class_labels,
+    guidance_scale=4.0,
+    num_inference_steps=80,
+    do_classifier_free_guidance=True,
+    device=None,
+):
+
+    batch_size = len(class_labels)
+    latent_size = pipe.transformer.config.sample_size
+    latent_channels = pipe.transformer.config.in_channels
+    
+    latents = start_latents.clone()
+    
+    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+
+    class_labels = torch.tensor(class_labels, device=device).reshape(-1)
+    class_null = torch.tensor([1000] * batch_size, device=device)
+    class_labels_input = torch.cat([class_labels, class_null], 0) if guidance_scale > 1 else class_labels
+    class_labels_input = class_labels_input.to(device)
+
+    # set step values
+    pipe.scheduler.set_timesteps(num_inference_steps)
+    intermediate_latents = []
+    
+    # Reversed timesteps <<<<<<<<<<<<<<<<<<<<
+    timesteps_all = reversed(pipe.scheduler.timesteps)
+    
+    for i in tqdm(range(1, num_inference_steps), total=num_inference_steps - 1):
+        
+        # We'll skip the final iteration
+        if i >= num_inference_steps - 1 -10:
+            continue
+
+        t = timesteps_all[i]    
+        
+        if do_classifier_free_guidance:
+            half = latent_model_input[: len(latent_model_input) // 2]
+            latent_model_input = torch.cat([half, half], dim=0)
+            
+        latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
+
+        timesteps = t
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = latent_model_input.device.type == "mps"
+            if isinstance(timesteps, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=latent_model_input.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(latent_model_input.device)
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(latent_model_input.shape[0])
+        latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
+        # predict noise model_output
+        noise_pred = pipe.transformer(
+            latent_model_input, timestep=timesteps, class_labels=class_labels_input
+        ).sample
+
+         # Perform guidance
+        if do_classifier_free_guidance:
+            # perform guidance
+            # noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            # noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        
+            eps, rest = noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
+            cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+
+            half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
+            eps = torch.cat([half_eps, half_eps], dim=0)
+
+            noise_pred = torch.cat([eps, rest], dim=1)
+
+        # learned sigma
+        if pipe.transformer.config.out_channels // 2 == latent_channels:
+        
+            model_output, _ = torch.split(noise_pred, latent_channels, dim=1)
+        else:
+            model_output = noise_pred
+
+        current_t = max(0, t.item() - (1000 // num_inference_steps))  # t
+        next_t = t  # min(999, t.item() + (1000//num_inference_steps)) # t+1
+        alpha_t = pipe.scheduler.alphas_cumprod[current_t]
+        alpha_t_next = pipe.scheduler.alphas_cumprod[next_t]
+
+        # Inverted update step (re-arranging the update step to get x(t) (new latents) as a function of x(t-1) (current latents)
+        latent_model_input = (latent_model_input - (1 - alpha_t).sqrt() * model_output) * (alpha_t_next.sqrt() / alpha_t.sqrt()) + (
+            1 - alpha_t_next
+        ).sqrt() * model_output
+
+        if guidance_scale > 1:
+            latents_out, _ = latent_model_input.chunk(2, dim=0)
+        else:
+            latents_out = latent_model_input
+
+        # Store i=[3, 8, 13, 18, 23, 28, 33, 38, 43, 48]
+        # if (i+2)%5 == 0:
+        if i>23 and i<39:
+            intermediate_latents.append(latents_out)
+    return torch.stack(intermediate_latents, dim=0)
+
+    # return torch.cat(intermediate_latents)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Script to train Stable Diffusion XL for InstructPix2Pix."
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="/scratch/zhao.lin1/ddim_inversion_statistic",
+        help="statistic save path",
+    )
+    parser.add_argument(
+        "--mapping_file",
+        type=str,
+        default="ds_inf/imagenet_1k_mapping.json",
+    )
+    parser.add_argument("--txt_file", default='ds_inf/imagenet1k_train.txt', type=str)
+    parser.add_argument("--pretrained_path", default='/scratch/zhao.lin1/DiT-XL-2-256', type=str)
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=200,
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=24,
+    )
+    parser.add_argument(
+        "--start",
+        type=int,
+        default=0,
+    )
+    parser.add_argument(
+        "--end",
+        type=int,
+        default=25,
+    )
+    parser.add_argument(
+        "--gpu",
+        type=int,
+        default=1,
+    )
+    args = parser.parse_args()
+    
+    return args
+
+def view_latents(pipe = None, inverted_latents = None):
+    with torch.no_grad():
+        im = pipe.decode_latents(inverted_latents[-1].unsqueeze(0))
+    pipe.numpy_to_pil(im)[0]
+    
+    
+def collate_fn(batch):
+    batch = [item for item in batch if item is not None]
+
+    if len(batch) == 0:
+        return None
+
+    images = torch.stack([item['images'] for item in batch])
+    labels = torch.tensor([item['labels'] for item in batch])
+    idx = torch.tensor([item['idx'] for item in batch])
+    paths = [item['paths'] for item in batch]
+
+    return {
+        'images': images,
+        'labels': labels,
+        'idx':idx,
+        'paths': paths
+    }
+
+
+def save_latent(latent, save_path):
+    torch.save(latent, save_path)
+
+
+
+class ImageNetDataset(Dataset):
+    def __init__(self, txt_file='', mapping_file=None, class_dir=None):
+        self.images = []
+        self.img_labels = []
+        self.class_dir = class_dir
+        self.transform = self.get_transforms()
+        
+        # Load class mapping and json file
+        self.wnid_to_index = load_mapping(mapping_file)
+        self._load_from_txt(txt_file)
+        
+
+    def _load_from_txt(self, txt_file):
+        with open(txt_file, "r") as file:
+            image_paths = file.readlines()
+        image_paths = [path.strip() for path in image_paths if path.split('/')[-2]==self.class_dir]
+        for path in image_paths:
+            self.images.append(path)
+            class_index = self.wnid_to_index[path.split('/')[-2]]
+            self.img_labels.append(class_index)
+
+
+    def get_transforms(self):
+        # normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+        #                             std=[0.229, 0.224, 0.225])
+        
+        return transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(256),
+            transforms.ToTensor(),
+            # normalize
+        ])
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        img_path = self.images[idx]
+        try:
+            image = Image.open(img_path).convert('RGB')
+        except Exception as e:
+            print(f"Error loading image {img_path}: {e}")
+            # Return a black image in case of error
+            image = Image.new('RGB', (256, 256))
+            
+        img_label = self.img_labels[idx]
+        
+        if self.transform:
+            image = self.transform(image)
+
+        sample = {
+            'images': image,
+            'paths': img_path,
+            'labels': img_label,
+            'idx': idx
+        }
+
+        return sample
+
+
+def load_mapping(mapping_file):
+    new_mapping = {}
+    with open(mapping_file, 'r') as file:
+        data = json.load(file)
+        if "tiny" in mapping_file:
+            for index, line in enumerate(file):
+                # Extract wnid (eg. n01443537) for each line and -1
+                key = line.split()[0]
+                new_mapping[key] = index 
+        else:
+            new_mapping = {item["wnid"]: item["index"] for item in data.values()}
+    return new_mapping
+
+def main():
+    args = parse_args()
+    torch.cuda.set_device(args.gpu)
+    device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")
+    
+    torch.cuda.set_device(args.gpu)
+    
+    wnid_to_index = load_mapping(args.mapping_file)
+    class_dirs = sorted(list(wnid_to_index.keys()))[args.start:args.end]
+ 
+    pipe = DiTPipeline.from_pretrained(args.pretrained_path, torch_dtype=torch.float16)
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) 
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    
+
+    pipe = pipe.to(device)
+
+    for class_dir in tqdm(class_dirs):
+        imgnet1k_dataset = ImageNetDataset(args.txt_file, args.mapping_file, class_dir)
+        trainloader = torch.utils.data.DataLoader(imgnet1k_dataset, batch_size=args.batch_size, shuffle=False,num_workers=args.num_workers, drop_last=False, collate_fn=collate_fn)
+        latents = []
+        for sample in tqdm(trainloader):
+            with torch.no_grad():
+                images = sample['images'].to(device)
+                latent = pipe.vae.encode(images.to(device, dtype=torch.float16) * 2 - 1)
+                ls = 0.18215 * latent.latent_dist.sample()
+                
+                inverted_latents = invert(pipe,start_latents = ls, class_labels=sample['labels'], num_inference_steps = 50, device=device).cpu()
+                latents.append(torch.flatten(inverted_latents.permute(1,0,2,3,4), start_dim=2))
+
+        latents = torch.cat(latents, dim=0).cpu()
+        mean = latents.mean(dim=0)
+        variance = latents.var(dim=0)
+        torch.save({"mean": mean, "variance": variance}, os.path.join(args.save_dir,class_dir+'.pt'))
+
+            
+
+
+
+if __name__ == "__main__":
+    main()

var/D3HR/generation/dit_inversion_save_statistic.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+
+save_dir='/home/v-qichen3/debug/var/D3HR/ddim_inversion_statistic'
+pretrained_path='/home/v-qichen3/debug/var/D3HR/DiT-XL'
+
+# the range of class ids
+# To improve efficiency, we distribute the generation of different classes across separate GPUs. You can change it to your own setting.
+n=0
+declare -a gpus=(0 1)
+declare -a starts=($n $(($n+100)))
+declare -a ends=($(($n+100)) $(($n+200)))
+
+for i in ${!gpus[@]}; do
+    gpu=${gpus[$i]}
+    start=${starts[$i]}
+    end=${ends[$i]}
+    
+    echo "Running on GPU $gpu with start=$start and end=$end"
+    python generation/dit_inversion_save_statistic.py --start $start --end $end --gpu $gpu --save_dir $save_dir --pretrained_path $pretrained_path &
+done
+
+# waiting for all tasks
+wait
+echo "All tasks completed."

var/D3HR/generation/group_sampling.py

@@ -0,0 +1,368 @@
+import torch
+import os
+import sys
+sys.path.append('/home/zhao.lin1/D3HR')
+from generation.dit_inversion_save_statistic import sample
+from diffusers import DiTPipeline, DDIMScheduler
+import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import argparse
+from tqdm import tqdm
+import ipdb
+from PIL import Image
+
+def save_average_latents(input_dir, save_dir):
+    latents_by_class = {}
+
+    # Iterate over each class folder in input_dir
+    for class_name in sorted(os.listdir(input_dir)):
+        class_dir = os.path.join(input_dir, class_name)
+        
+        if os.path.isdir(class_dir):
+            latents = []
+            for file_name in sorted(os.listdir(class_dir)):
+                if file_name.endswith('.pt'):
+                    file_path = os.path.join(class_dir, file_name)
+                    latent = torch.load(file_path)  # Load the latent .pth file
+                    latents.append(latent)
+            
+            # Store the latent vectors for this class
+            latents_tensor = torch.stack(latents)
+            average_latent = torch.mean(latents_tensor, dim=0)
+
+            # latents_by_class[class_name] = latents
+
+            save_file_path = os.path.join(save_dir, f'{class_name}_average_latent.pth')
+            torch.save(average_latent, save_file_path)
+            print(f"Saved average latent for class '{class_name}' to {save_file_path}")
+    return 
+
+def load_mapping(mapping_file):
+    new_mapping = {}
+    with open(mapping_file, 'r') as file:
+        data = json.load(file)
+        if "tiny" in mapping_file:
+            for index, line in enumerate(file):
+                #Extract the wnid starting with 'n' from each line and subtract 1 from the line number.
+                key = line.split()[0]
+                new_mapping[key] = index 
+        else:
+            new_mapping = {item["wnid"]: item["index"] for item in data.values()}
+    return new_mapping
+
+
+def process_class(class_dir, folder_path):
+    data_list = []
+    class_path = os.path.join(folder_path, class_dir)
+    for file_name in os.listdir(class_path):
+        if file_name.endswith(".pt"):
+            file_path = os.path.join(class_path, file_name)
+            tensor = torch.load(file_path, map_location=torch.device('cpu'))
+            data_list.append(torch.flatten(tensor, start_dim=1))
+
+    if data_list: 
+        data = torch.stack(data_list, dim=0)
+        mean = data.mean(dim=0)
+        variance = data.var(dim=0)
+        return class_dir, {"mean": mean, "variance": variance}
+    return class_dir, None  
+
+def process_p_sample(class_dir, folder_path):
+    data_list = []
+    class_path = os.path.join(folder_path, class_dir)
+    for file_name in os.listdir(class_path):
+        if file_name.endswith(".pt"):
+            file_path = os.path.join(class_path, file_name)
+            tensor = torch.load(file_path, map_location=torch.device('cpu'))
+            data_list.append(tensor.flatten())
+
+    data = torch.vstack(data_list)
+    return data
+
+
+def kl_divergence(selected_points, mean, std, device):
+    """
+    Compute the KL divergence between the candidate distribution and the target Gaussian distribution.
+    KL(P || Q) = 0.5 * [tr(Sigma_Q^-1 Sigma_P) + (mu_Q - mu_P)^T Sigma_Q^-1 (mu_Q - mu_P) - k + log(det(Sigma_Q) / det(Sigma_P))]
+    Here, the target distribution Q is N(mean,cov), 
+    and the sampling distribution P is estimated from the selected_points.
+    """
+    # k = mean.size(0)  # Feature dimension: 4090
+    selected_mean = selected_points.mean(dim=0)
+    selected_var = selected_points.var(dim=0)
+    selected_std = torch.sqrt(selected_var)
+
+    # Compute KL divergnece
+    diff = mean - selected_mean
+    log_sigma_ratio = torch.log(selected_std / std)
+    variance_ratio = (std**2 + diff**2) / (2 * selected_std**2)
+    kl = torch.sum(log_sigma_ratio + variance_ratio - 0.5)
+
+    return kl.item()  # Return a scalar value
+
+def kl_divergence_independent_batch(mean, std, samples, device):
+    mean = mean.to(device)
+    std = std.to(device)
+    samples = samples.to(device)
+
+    # Compute KL divergence
+    diff = samples - mean
+    term1 = torch.sum(diff**2, dim=1) / (2 * std**2)  
+    kl_divs = term1 + torch.log(std) - 0.5            # log(std/std)=0
+    return kl_divs
+
+def sinkhorn(A, B, epsilon=0.1, max_iter=1000, tol=1e-9):
+    """
+    Estimate the Wasserstein distance using the Sinkhorn algorithm, which supports distributions with different numbers of samples.
+    A, B: The two input distributions
+    epsilon: Sinkhorn regularization parameter
+    max_iter: Maximum number of iterations
+    tol: Convergence tolerance
+    """
+
+    # The amount of samples
+    n_a, n_b = A.size(0), B.size(0)
+    
+    # Define weights and ensure normalization.
+    weight_a = torch.ones(n_a, device=A.device) / n_a
+    weight_b = torch.ones(n_b, device=B.device) / n_b
+
+    # Compute the distance matrix
+    C = torch.cdist(A, B, p=2) ** 2  # Squared Euclidean distance
+
+    # Initialize dual variables
+    u = torch.zeros(n_a, device=A.device)
+    v = torch.zeros(n_b, device=B.device)
+
+    K = torch.exp(-C / epsilon)  # Regularized distance matrix
+
+    for _ in range(max_iter):
+        # Update u and c, consider weights simultaneously
+        u_new = epsilon * torch.log(weight_a) - epsilon * torch.logsumexp(-K / epsilon + v.view(1, -1), dim=1)
+        v_new = epsilon * torch.log(weight_b) - epsilon * torch.logsumexp(-K / epsilon + u_new.view(-1, 1), dim=0)
+
+        # Check convergence
+        if torch.max(torch.abs(u_new - u)) < tol and torch.max(torch.abs(v_new - v)) < tol:
+            break
+
+        u, v = u_new, v_new
+
+    transport_cost = torch.sum(K * C)  
+    wasserstein_distance = transport_cost + epsilon * (torch.sum(u * weight_a) + torch.sum(v * weight_b))
+
+    return wasserstein_distance
+
+def skewness(tensor):
+    mean = torch.mean(tensor, dim=0)
+    std = torch.std(tensor, dim=0)
+    n = tensor.size(0)
+    skew = torch.sum(((tensor - mean) / std) ** 3, dim=0) * (n / ((n - 1) * (n - 2)))
+    return skew
+
+def kurtosis(tensor):
+    mean = torch.mean(tensor, dim=0)
+    std = torch.std(tensor, dim=0)
+    n = tensor.size(0)
+    kurt = torch.sum(((tensor - mean) / std) ** 4, dim=0) * (n * (n + 1)) / ((n - 1) * (n - 2) * (n - 3)) - (3 * (n - 1) ** 2) / ((n - 2) * (n - 3))
+    return kurt
+
+
+def skewness_batch(tensor):
+    mean = torch.mean(tensor, dim=1, keepdim=True)  # shape: [20000, 1, 4096]
+    std = torch.std(tensor, dim=1, keepdim=True)    # shape: [20000, 1, 4096]
+    
+    n = tensor.size(2)  # feature dimension: 4096
+    skew = torch.sum(((tensor - mean) / std) ** 3, dim=2) * (n / ((n - 1) * (n - 2)))  # shape: [20000, 1]
+    
+    return skew
+
+
+def evaluate_distribution(samples, mean, std):
+    sample_mean = torch.mean(samples, dim=0)
+    mean_diff = torch.norm(sample_mean - mean)
+
+    sample_std = torch.std(samples, dim=0)
+    std_diff = torch.norm(sample_std - std)
+
+    sample_skew = skewness_batch(samples)
+
+    skew_diff = torch.norm(torch.tensor(sample_skew) - 0)  # Sample Skewness close to 0
+    # kurt_diff = torch.norm(torch.tensor(sample_kurt) - 3)  # Sample Kurtosis close to 3
+
+    # Comprehensive evaluation: each component can be weighted as needed
+    score = mean_diff + std_diff + 10*skew_diff
+    return score
+
+def select_algorithm(n_trials, n_samples, mean, std, device):
+    best_score = float('inf')
+    best_sample = None
+    for _ in range(n_trials):
+        samples = torch.normal(mean.expand(n_samples, -1), std.expand(n_samples, -1)).to(device)
+        
+        score = evaluate_distribution(samples, mean, std)
+        print(score)
+
+        # Choose sample with best score
+        if score < best_score:
+            best_score = score
+            best_sample = samples
+
+    return best_sample
+
+
+def evaluate_distribution_batch(samples, mean, std):
+    sample_mean = torch.mean(samples, dim=1) 
+    mean_diff = torch.norm(sample_mean - mean, dim=1) 
+
+    sample_std = torch.std(samples, dim=1)  
+    std_diff = torch.norm(sample_std - std, dim=1) 
+
+    # Compute samples Skewness
+    sample_skew = skewness_batch(samples)
+    # Batch computation of skewness differences, default is 0 -> sample_skew-0
+    skew_diff = torch.norm(sample_skew, dim=1)  
+
+    # Comprehensive evaluation: each component can be weighted as needed
+    score = mean_diff + std_diff + 0.1 * skew_diff
+    return score
+
+def select_algorithm_batch(n_trials, n_samples, mean, std, device, seed):
+    if seed is not None:
+        torch.manual_seed(seed)
+    # Batch computation, where n_trials indicates the batch size.
+    samples = torch.normal(mean.expand(n_trials, n_samples, -1), std.expand(n_trials, n_samples, -1)).to(device)  # Batch sampling
+    scores = evaluate_distribution_batch(samples, mean, std)  # Batch evaluating
+    
+    best_score, best_idx = torch.min(scores, dim=0)  # Find smaples with best (small) scores
+    worst_score, worst_idx = torch.max(scores, dim=0)
+    best_sample = samples[best_idx]  # Get best samples
+    worst_sample = samples[worst_idx]
+    
+    return best_sample, worst_sample, best_score, worst_score
+
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Script to train Stable Diffusion XL for InstructPix2Pix."
+    )
+    parser.add_argument(
+        "--mapping_file",
+        type=str,
+        default="ds_inf/imagenet_1k_mapping.json",
+    )
+    parser.add_argument("--pretrained_path", default='/scratch/zhao.lin1/DiT-XL-2-256', type=str)
+    parser.add_argument("--save_dir", default='/scratch/zhao.lin1/distilled_images/', type=str)
+    parser.add_argument("--statistic_path", default='/scratch/zhao.lin1/ddim_inversion_statistic', type=str)
+    parser.add_argument(
+        "--start",
+        type=int,
+        default=0,
+    )
+    parser.add_argument(
+        "--end",
+        type=int,
+        default=1000,
+    )
+    parser.add_argument(
+        "--gpu",
+        type=int,
+        default=0,
+    )
+    parser.add_argument(
+        "--ipc",
+        type=int,
+        default=20,
+    )
+    parser.add_argument(
+        "--start_step",
+        type=int,
+        default=18,
+    )
+    parser.add_argument(
+        "--i_step",
+        type=int,
+        default=6,
+    )
+    parser.add_argument(
+        "--m",
+        type=int,
+        default=100000,
+    )
+    args = parser.parse_args()
+    
+    return args
+
+
+def main():
+    args = parse_args()
+
+    device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")
+    
+    torch.cuda.set_device(device)
+    
+    # x5_step i=[3, 8, 13, 18, 23, 28, 33, 38, 43, 48] start_step = [45, 40, 35, 30, 25, 20, 15, 10, 5, 0]
+    # 10-20_step i=[24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39] start_step = [24, 23, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9]
+
+    wnid_to_index = load_mapping(args.mapping_file)
+    class_dirs = sorted(list(wnid_to_index.keys()))[args.start:args.end]
+
+
+
+    pipe = DiTPipeline.from_pretrained(args.pretrained_path, torch_dtype=torch.float16)
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) 
+    pipe = pipe.to(device)
+
+    for class_dir in tqdm(class_dirs):
+        # Compute the best and worst samples
+        statics = torch.load(os.path.join(args.statistic_path, class_dir+'.pt'))
+        mean = statics['mean'][args.i_step].to(device)
+        variance = statics['variance'][args.i_step].to(device)
+        std = torch.sqrt(variance)
+
+        latents_best = None
+        latents_worst = None
+        best_overall_score = float('inf')  # initialize as inf
+        worst_overall_score = float('-inf')  # initialize as -ing
+
+        # group sampling
+        for i in range(args.m//10):
+            seed = i * 12345
+            best_sample, worst_sample, best_score, worst_score = select_algorithm_batch(10000, args.ipc, mean, std, device, seed)
+            
+            # Update best and worst samples
+            if best_score < best_overall_score:
+                best_overall_score = best_score
+                latents_best = best_sample
+            
+            if worst_score > worst_overall_score:
+                worst_overall_score = worst_score
+                latents_worst = worst_sample
+
+        # Output results
+        print("Best overall score:", best_overall_score)
+        print("Worst overall score:", worst_overall_score)
+        
+
+        latents_best = latents_best.view(-1,4,32,32)
+        # latents_worst = latents_worst.view(-1,4,32,32)
+
+
+        # Generate images
+        for k, latent in enumerate(latents_best):            
+            image = sample(
+                pipe,
+                class_labels=torch.tensor(wnid_to_index[class_dir]).unsqueeze(0),
+                start_latents=latent.unsqueeze(0).to(torch.float16),
+                start_step=args.start_step,
+                num_inference_steps=50,
+                device=device
+            )
+            os.makedirs(os.path.join(args.save_dir,class_dir.split('/')[-1]), exist_ok=True)
+            pipe.numpy_to_pil(image)[0].resize((224, 224), Image.LANCZOS).save(os.path.join(args.save_dir,class_dir.split('/')[-1], str(k)+'.png'))
+            
+
+
+if __name__ == "__main__":
+    main()
+

var/D3HR/generation/group_sampling.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# Define the range of classes
+# To improve efficiency, we distribute the generation of different classes across separate GPUs. You can change it to your own setting.
+n=0
+declare -a gpus=(0 1)
+declare -a starts=($n $(($n+100)))
+declare -a ends=($(($n+100)) $(($n+200)))
+
+for i in ${!gpus[@]}; do
+    gpu=${gpus[$i]}
+    start=${starts[$i]}
+    end=${ends[$i]}
+    
+    echo "Running on GPU $gpu with start=$start and end=$end"
+    python generation/group_sampling.py --start $start --end $end --gpu $gpu &
+done
+
+wait
+echo "All tasks completed."

var/D3HR/imgs/framework.pdf

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

var/D3HR/requirements.txt

@@ -0,0 +1,9 @@
+torch==2.6.0
+torchvision==0.21.0
+transformers==4.46.2
+diffusers==0.28.0
+matplotlib==3.10.1
+ipdb==0.13.13
+scipy==1.15.2
+huggingface-hub==0.30.2
+accelerate==1.3.0

var/D3HR/validation/argument.py

@@ -0,0 +1,310 @@
+import os
+import argparse
+
+
+def str2bool(v):
+    """Cast string to boolean
+    """
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+parser = argparse.ArgumentParser("EEF")
+parser.add_argument(
+    "--arch-name",
+    type=str,
+    default="resnet18",
+    help="arch name from pretrained torchvision models",
+)
+parser.add_argument(
+    "--subset",
+    type=str,
+    default="imagenet-1k",
+)
+parser.add_argument(
+    "--spec",
+    type=str,
+    default="none",
+)
+parser.add_argument(
+    "--label-smoothing", default=0.0, type=float, help="label smoothing (default: 0.0)", dest="label_smoothing"
+)
+parser.add_argument(
+    "--data-dir",
+    nargs='+',
+    default=["../data/imagenet"],
+    help="path to imagenet dataset",
+)
+parser.add_argument(
+    "--nclass",
+    type=int,
+    default=10,
+    help="number of classes for synthesis or validation",
+)
+parser.add_argument(
+    "--ipc",
+    type=int,
+    default=10,
+    help="number of images per class for synthesis or validation",
+)
+parser.add_argument(
+    "--target-ipc",
+    type=int,
+    default=50,
+    help="number of images per class for synthesis or validation",
+)
+parser.add_argument(
+    "--phase",
+    type=int,
+    default=0,
+)
+parser.add_argument(
+    "--input-size",
+    default=224,
+    type=int,
+    metavar="S",
+)
+parser.add_argument(
+    "--save-size",
+    default=224,
+    type=int,
+    metavar="S",
+)
+parser.add_argument(
+    "--repeat",
+    default=1,
+    type=int,
+    help="Repeat times for the validation"
+)
+parser.add_argument(
+    "--factor",
+    default=2,
+    type=int,
+)
+parser.add_argument(
+    "--batch-size", default=0, type=int, metavar="N"
+)
+parser.add_argument(
+    "--accum-steps",
+    type=int,
+    default=1,
+    help="gradient accumulation steps for small gpu memory",
+)
+parser.add_argument(
+    "--mix-type",
+    default="cutmix",
+    type=str,
+    choices=["mixup", "cutmix", None],
+    help="mixup or cutmix or None",
+)
+parser.add_argument(
+    "--stud-name",
+    type=str,
+    default="resnet18",
+    help="arch name from torchvision models",
+)
+parser.add_argument(
+    "--workers",
+    default=24,
+    type=int,
+    metavar="N",
+    help="number of data loading workers (default: 4)",
+)
+parser.add_argument(
+    "--temperature",
+    type=float,
+    help="temperature for distillation loss",
+)
+parser.add_argument(
+    "--min-scale-crops", type=float, default=0.08, help="argument in RandomResizedCrop"
+)
+parser.add_argument(
+    "--max-scale-crops", type=float, default=1, help="argument in RandomResizedCrop"
+)
+parser.add_argument("--epochs", default=300, type=int)
+parser.add_argument(
+    "--results-dir",
+    type=str,
+    default="results",
+    help="where to store synthetic data",
+)
+parser.add_argument(
+    "--seed", default=42, type=int, help="seed for initializing training. "
+)
+parser.add_argument(
+    "--mixup",
+    type=float,
+    default=0.8,
+    help="mixup alpha, mixup enabled if > 0. (default: 0.8)",
+)
+parser.add_argument(
+    "--cutmix",
+    type=float,
+    default=1.0,
+    help="cutmix alpha, cutmix enabled if > 0. (default: 1.0)",
+)
+parser.add_argument("--cos", default=True, help="cosine lr scheduler")
+parser.add_argument("--verbose", type=str2bool, default=False)
+parser.add_argument("--mapping_file", default="ds_inf/imagenet_1k_mapping.json", type=str)
+parser.add_argument("--txt_file", default='/home/zhao.lin1/DD-DDIM-inversion/ds_inf/imagenet-1k/biggest_20%_ipc_for_all_1k.txt', type=str)
+parser.add_argument("--val_txt_file", default='/home/zhao.lin1/CONCORD/val.txt', type=str)
+# diffusion
+parser.add_argument("--dit-model", default='DiT-XL/2')
+parser.add_argument("--ckpt", type=str, default='pretrained_models/DiT-XL-2-256x256.pt',
+                    help="Optional path to a DiT checkpoint (default: auto-download a pre-trained DiT-XL/2 model).")
+parser.add_argument("--dit-image-size", default=256, type=int)
+parser.add_argument("--num-dit-classes", default=1000, type=int)
+parser.add_argument("--diffusion-steps", default=1000, type=int)
+parser.add_argument("--cfg-scale", type=float, default=4.0)
+
+parser.add_argument("--vae-path", default='stabilityai/sd-vae-ft-ema')
+
+# distillation
+parser.add_argument("--save-path", default='./results/test')
+parser.add_argument("--description-path", default='./misc/class_description.json')
+parser.add_argument("--clip-alpha", type=float, default=10.0)
+parser.add_argument("--cls-alpha", type=float, default=10.0)
+parser.add_argument("--num-neg-samples", type=int, default=5)
+parser.add_argument("--neg-policy", type=str, default="weighted")
+
+# sgd
+parser.add_argument("--sgd", default=False, action="store_true", help="sgd optimizer")
+parser.add_argument(
+    "-lr",
+    "--learning-rate",
+    type=float,
+    default=0.1,
+    help="sgd init learning rate",
+)
+parser.add_argument("--momentum", type=float, default=0.9, help="sgd momentum")
+parser.add_argument("--weight-decay", type=float, default=1e-4, help="sgd weight decay")
+
+# adamw
+parser.add_argument("--adamw-lr", type=float, default=0, help="adamw learning rate")
+parser.add_argument(
+    "--adamw-weight-decay", type=float, default=0.01, help="adamw weight decay"
+)
+parser.add_argument(
+    "--exp-name",
+    type=str,
+    help="name of the experiment, subfolder under syn_data_path",
+)
+args = parser.parse_args()
+
+# temperature
+if args.mix_type == "mixup":
+    args.temperature = 4
+elif args.mix_type == "cutmix":
+    args.temperature = 20
+
+if args.subset == "imagenet_1k":
+    args.nclass = 1000
+    args.classes = range(args.nclass)
+    args.val_ipc = 50
+    args.input_size = 224
+
+elif args.subset == "imagewoof":
+    args.nclass = 10
+    args.classes = range(args.nclass)
+    args.val_ipc = 50
+    args.input_size = 224
+    if args.ipc == 10:
+        args.epochs = 2000
+    elif args.ipc == 50:
+        args.epochs = 1500
+    else:
+        args.epochs = 1000
+
+
+elif args.subset == "cifar10":
+    args.nclass = 10
+    args.classes = range(args.nclass)
+    args.val_ipc = 1000
+    args.input_size = 32
+    args.epochs = 1000
+
+elif args.subset == "cifar100":
+    args.nclass = 100
+    args.classes = range(args.nclass)
+    args.val_ipc = 100
+    args.input_size = 32
+    args.epochs = 400
+
+elif args.subset == "tinyimagenet":
+    args.nclass = 200
+    args.classes = range(args.nclass)
+    args.val_ipc = 50
+    args.input_size = 64
+    args.epochs = 300
+
+# set up batch size
+if args.batch_size == 0:
+    if args.ipc >= 50:
+        args.batch_size = 100
+    elif args.ipc >= 10:
+        args.batch_size = 50
+    elif args.ipc > 0:
+        args.batch_size = 15
+    elif args.ipc == -1:
+        args.batch_size = 100
+
+    if args.nclass == 10:
+        args.batch_size *= 1
+    if args.nclass == 100:
+        args.batch_size *= 2
+    # if args.nclass == 1000:
+    #     args.batch_size *= 2
+
+# reset batch size below ipc * nclass
+if args.ipc != -1 and args.batch_size > args.ipc * args.nclass:
+    args.batch_size = int(args.ipc * args.nclass)
+
+# reset batch size with accum_steps
+if args.accum_steps != 1:
+    args.batch_size = int(args.batch_size / args.accum_steps)
+
+# result dir for saving
+args.exp_name = f"{args.spec}_{args.arch_name}_f{args.factor}_ipc{args.ipc}"
+if not os.path.exists(f"./exp/{args.exp_name}"):
+    os.makedirs(f"./exp/{args.exp_name}")
+
+
+# adamw learning rate
+if args.stud_name == "vgg11":
+    args.adamw_lr = 0.0005
+elif args.stud_name == "conv3":
+    args.adamw_lr = 0.001
+elif args.stud_name == "conv4":
+    args.adamw_lr = 0.001
+elif args.stud_name == "conv5":
+    args.adamw_lr = 0.001
+elif args.stud_name == "conv6":
+    args.adamw_lr = 0.001
+elif args.stud_name == "resnet18":
+    args.adamw_lr = 0.001
+elif args.stud_name == "resnet18_modified":
+    args.adamw_lr = 0.001
+elif args.stud_name == "efficientnet_b0":
+    args.adamw_lr = 0.002
+elif args.stud_name == "mobilenet_v2":
+    args.adamw_lr = 0.0025
+elif args.stud_name == "alexnet":
+    args.adamw_lr = 0.0001
+elif args.stud_name == "resnet50":
+    args.adamw_lr = 0.001
+elif args.stud_name == "resnet50_modified":
+    args.adamw_lr = 0.001
+elif args.stud_name == "resnet101":
+    args.adamw_lr = 0.001
+elif args.stud_name == "resnet101_modified":
+    args.adamw_lr = 0.001
+elif args.stud_name == "vit_b_16":
+    args.adamw_lr = 0.0001
+elif args.stud_name == "swin_v2_t":
+    args.adamw_lr = 0.0001

var/D3HR/validation/get_train_list.py

@@ -0,0 +1,26 @@
+import json
+import os
+
+def load_mapping(mapping_file):
+    new_mapping = {}
+    with open(mapping_file, 'r') as file:
+        data = json.load(file)
+        if "tiny" in mapping_file:
+            for index, line in enumerate(file):
+                key = line.split()[0]
+                new_mapping[key] = index 
+        else:
+            new_mapping = {item["wnid"]: item["index"] for item in data.values()}
+    return new_mapping
+
+
+wnid_to_index = load_mapping("ds_inf/imagenet_1k_mapping.json")
+class_dirs = sorted(list(wnid_to_index.keys()))
+path_list = []
+for class_dir in class_dirs:
+    for i in range(20):
+        path_list.append(os.path.join('/scratch/zhao.lin1/imagenet1k_256_4.0classfree_start_step_18_ddim_inversion_20_min_images_2/', class_dir, str(i)+'.png'))
+output_file = "/scratch/zhao.lin1/imagenet1k_256_4.0classfree_start_step_18_ddim_inversion_20_min_images_2/train.txt"
+with open(output_file, "w") as file:
+    for path in path_list:
+        file.write(path + "\n") 

var/D3HR/validation/models/__init__.py

@@ -0,0 +1,190 @@
+import torch
+import torch.nn as nn
+import torchvision.models as thmodels
+
+from .convnet import ConvNet
+from .resnet import resnet18, resnet50, resnet101, resnet152
+from .mobilenet_v2 import mobilenetv2
+# import timm
+
+
+
+def load_model(model_name="resnet18", dataset="cifar10", spec='full', pretrained=True, input_size=224, classes=[]):
+    def get_model(model_name="resnet18"):
+        if "conv" in model_name:
+            size = input_size
+            nclass = 1000
+
+            model = ConvNet(
+                num_classes=nclass,
+                net_norm="batch",
+                net_act="relu",
+                net_pooling="avgpooling",
+                net_depth=int(model_name[-1]),
+                net_width=128,
+                channel=3,
+                im_size=(size, size),
+            )
+        elif model_name == 'resnet18':
+            model = resnet18(weights=None)
+        elif model_name == 'resnet50':
+            model = resnet50(weights=None)
+        elif model_name == 'resnet101':
+            model = resnet101(weights=None)
+        elif model_name == 'resnet152':
+            model = resnet152(weights=None)
+        elif model_name == 'mobilenet_v2':
+            model = mobilenetv2()
+        elif model_name == 'efficientnet_b0':
+            model = timm.create_model('efficientnet_b0.ra_in1k', pretrained=False)
+        elif model_name == "resnet18_modified":
+            model = thmodels.__dict__["resnet18"](pretrained=False)
+            model.conv1 = nn.Conv2d(
+                3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
+            )
+            model.maxpool = nn.Identity()
+        elif model_name == "resnet50_modified":
+            model = thmodels.__dict__["resnet50"](pretrained=False)
+            model.conv1 = nn.Conv2d(
+                3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
+            )
+            model.maxpool = nn.Identity()
+        elif model_name == "resnet101_modified":
+            model = thmodels.__dict__["resnet101"](pretrained=False)
+            model.conv1 = nn.Conv2d(
+                3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
+            )
+            model.maxpool = nn.Identity()
+        else:
+            model = thmodels.__dict__[model_name](weights=None)
+
+        return model
+
+    def pruning_classifier(model=None, classes=[]):
+        try:
+            model_named_parameters = [name for name, x in model.named_parameters()]
+            for name, x in model.named_parameters():
+                if (
+                    name == model_named_parameters[-1]
+                    or name == model_named_parameters[-2]
+                ):
+                    x.data = x[classes]
+        except:
+            print("ERROR in changing the number of classes.")
+
+        return model
+
+    model = get_model(model_name)
+    model = pruning_classifier(model, classes)
+
+    if pretrained:
+        if dataset == 'imagenet_1k':
+            if model_name == "efficientnet_b0":
+                checkpoint = timm.create_model('efficientnet_b0.ra_in1k', pretrained=True).state_dict()
+                model.load_state_dict(checkpoint)
+            elif model_name == 'conv4':
+                state_dict = torch.load('/home/linz/CONCORD/pretrained_models/imagenet-1k_conv4.pth')
+                model.load_state_dict(state_dict['model'])
+            elif model_name == 'resnet18':
+                model = resnet18(weights='DEFAULT')
+            elif model_name == 'mobilenet_v2':
+                model.load_state_dict(torch.load('/home/zhao.lin1/CONCORD/pretrained_models/mobilenetv2_1.0-0c6065bc.pth'))
+            else:
+                raise AttributeError(f'{model_name} is not supported in the pre-trained pool')
+        else:
+            checkpoint = torch.load(
+                f"pretrain_models/{dataset}_{model_name}.pth", map_location="cpu"
+            )
+            model.load_state_dict(checkpoint["model"])
+
+
+    return model
+
+
+# def load_model(model_name="resnet18", dataset="cifar10", pretrained=True, classes=[]):
+#     def get_model(model_name="resnet18"):
+#         if "conv" in model_name:
+#             if dataset in ["cifar10", "cifar100"]:
+#                 size = 32
+#             elif dataset == "tinyimagenet":
+#                 size = 64
+#             elif dataset in ["imagenet-nette", "imagenet-woof", "imagenet-100"]:
+#                 size = 128
+#             else:
+#                 size = 224
+
+#             nclass = len(classes)
+
+#             model = ConvNet(
+#                 num_classes=nclass,
+#                 net_norm="batch",
+#                 net_act="relu",
+#                 net_pooling="avgpooling",
+#                 net_depth=int(model_name[-1]),
+#                 net_width=128,
+#                 channel=3,
+#                 im_size=(size, size),
+#             )
+#         elif model_name == "resnet18_modified":
+#             model = thmodels.__dict__["resnet18"](pretrained=False)
+#             model.conv1 = nn.Conv2d(
+#                 3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
+#             )
+#             model.maxpool = nn.Identity()
+#         elif model_name == "resnet101_modified":
+#             model = thmodels.__dict__["resnet101"](pretrained=False)
+#             model.conv1 = nn.Conv2d(
+#                 3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
+#             )
+#             model.maxpool = nn.Identity()
+#         else:
+#             model = thmodels.__dict__[model_name](pretrained=False)
+
+#         return model
+
+#     def pruning_classifier(model=None, classes=[]):
+#         try:
+#             model_named_parameters = [name for name, x in model.named_parameters()]
+#             for name, x in model.named_parameters():
+#                 if (
+#                     name == model_named_parameters[-1]
+#                     or name == model_named_parameters[-2]
+#                 ):
+#                     x.data = x[classes]
+#         except:
+#             print("ERROR in changing the number of classes.")
+
+#         return model
+
+#     # "imagenet-100" "imagenet-10" "imagenet-first" "imagenet-nette" "imagenet-woof"
+#     model = get_model(model_name)
+#     model = pruning_classifier(model, classes)
+#     if pretrained:
+#         if dataset in [
+#             "imagenet-100",
+#             "imagenet-10",
+#             "imagenet-nette",
+#             "imagenet-woof",
+#             "tinyimagenet",
+#             "cifar10",
+#             "cifar100",
+#         ]:
+#             checkpoint = torch.load(
+#                 f"./data/pretrain_models/{dataset}_{model_name}.pth", map_location="cpu"
+#             )
+#             model.load_state_dict(checkpoint["model"])
+#         elif dataset in ["imagenet-1k"]:
+#             if model_name == "efficientNet-b0":
+#                 # Specifically, for loading the pre-trained EfficientNet model, the following modifications are made
+#                 from torchvision.models._api import WeightsEnum
+#                 from torch.hub import load_state_dict_from_url
+
+#                 def get_state_dict(self, *args, **kwargs):
+#                     kwargs.pop("check_hash")
+#                     return load_state_dict_from_url(self.url, *args, **kwargs)
+
+#                 WeightsEnum.get_state_dict = get_state_dict
+
+#             model = thmodels.__dict__[model_name](pretrained=True)
+
+#     return model

var/D3HR/validation/models/convnet.py

@@ -0,0 +1,147 @@
+import torch
+import torch.nn as nn
+
+
+# Conv-3 model
+class ConvNet(nn.Module):
+    def __init__(
+        self,
+        num_classes,
+        net_norm="batch",
+        net_depth=3,
+        net_width=128,
+        channel=3,
+        net_act="relu",
+        net_pooling="avgpooling",
+        im_size=(32, 32),
+    ):
+        # print(f"Define Convnet (depth {net_depth}, width {net_width}, norm {net_norm})")
+        super(ConvNet, self).__init__()
+        if net_act == "sigmoid":
+            self.net_act = nn.Sigmoid()
+        elif net_act == "relu":
+            self.net_act = nn.ReLU()
+        elif net_act == "leakyrelu":
+            self.net_act = nn.LeakyReLU(negative_slope=0.01)
+        else:
+            exit("unknown activation function: %s" % net_act)
+
+        if net_pooling == "maxpooling":
+            self.net_pooling = nn.MaxPool2d(kernel_size=2, stride=2)
+        elif net_pooling == "avgpooling":
+            self.net_pooling = nn.AvgPool2d(kernel_size=2, stride=2)
+        elif net_pooling == "none":
+            self.net_pooling = None
+        else:
+            exit("unknown net_pooling: %s" % net_pooling)
+
+        self.depth = net_depth
+        self.net_norm = net_norm
+
+        self.layers, shape_feat = self._make_layers(
+            channel, net_width, net_depth, net_norm, net_pooling, im_size
+        )
+        num_feat = shape_feat[0] * shape_feat[1] * shape_feat[2]
+        self.classifier = nn.Linear(num_feat, num_classes)
+
+    def forward(self, x, return_features=False):
+        for d in range(self.depth):
+            x = self.layers["conv"][d](x)
+            if len(self.layers["norm"]) > 0:
+                x = self.layers["norm"][d](x)
+            x = self.layers["act"][d](x)
+            if len(self.layers["pool"]) > 0:
+                x = self.layers["pool"][d](x)
+
+        # x = nn.functional.avg_pool2d(x, x.shape[-1])
+        out = x.view(x.shape[0], -1)
+        logit = self.classifier(out)
+
+        if return_features:
+            return logit, out
+        else:
+            return logit
+
+    def get_feature(
+        self, x, idx_from, idx_to=-1, return_prob=False, return_logit=False
+    ):
+        if idx_to == -1:
+            idx_to = idx_from
+        features = []
+
+        for d in range(self.depth):
+            x = self.layers["conv"][d](x)
+            if self.net_norm:
+                x = self.layers["norm"][d](x)
+            x = self.layers["act"][d](x)
+            if self.net_pooling:
+                x = self.layers["pool"][d](x)
+            features.append(x)
+            if idx_to < len(features):
+                return features[idx_from : idx_to + 1]
+
+        if return_prob:
+            out = x.view(x.size(0), -1)
+            logit = self.classifier(out)
+            prob = torch.softmax(logit, dim=-1)
+            return features, prob
+        elif return_logit:
+            out = x.view(x.size(0), -1)
+            logit = self.classifier(out)
+            return features, logit
+        else:
+            return features[idx_from : idx_to + 1]
+
+    def _get_normlayer(self, net_norm, shape_feat):
+        # shape_feat = (c * h * w)
+        if net_norm == "batch":
+            norm = nn.BatchNorm2d(shape_feat[0], affine=True)
+        elif net_norm == "layer":
+            norm = nn.LayerNorm(shape_feat, elementwise_affine=True)
+        elif net_norm == "instance":
+            norm = nn.GroupNorm(shape_feat[0], shape_feat[0], affine=True)
+        elif net_norm == "group":
+            norm = nn.GroupNorm(4, shape_feat[0], affine=True)
+        elif net_norm == "none":
+            norm = None
+        else:
+            norm = None
+            exit("unknown net_norm: %s" % net_norm)
+        return norm
+
+    def _make_layers(
+        self, channel, net_width, net_depth, net_norm, net_pooling, im_size
+    ):
+        layers = {"conv": [], "norm": [], "act": [], "pool": []}
+
+        in_channels = channel
+        if im_size[0] == 28:
+            im_size = (32, 32)
+        shape_feat = [in_channels, im_size[0], im_size[1]]
+
+        for d in range(net_depth):
+            layers["conv"] += [
+                nn.Conv2d(
+                    in_channels,
+                    net_width,
+                    kernel_size=3,
+                    padding=3 if channel == 1 and d == 0 else 1,
+                )
+            ]
+            shape_feat[0] = net_width
+            if net_norm != "none":
+                layers["norm"] += [self._get_normlayer(net_norm, shape_feat)]
+            layers["act"] += [self.net_act]
+            in_channels = net_width
+            if net_pooling != "none":
+                layers["pool"] += [self.net_pooling]
+                shape_feat[1] //= 2
+                shape_feat[2] //= 2
+
+        layers["conv"] = nn.ModuleList(layers["conv"])
+        layers["norm"] = nn.ModuleList(layers["norm"])
+        layers["act"] = nn.ModuleList(layers["act"])
+        layers["pool"] = nn.ModuleList(layers["pool"])
+        layers = nn.ModuleDict(layers)
+
+        return layers, shape_feat

var/D3HR/validation/models/dit_models.py

@@ -0,0 +1,438 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# GLIDE: https://github.com/openai/glide-text2im
+# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.jit import Final
+import numpy as np
+import math
+from timm.models.vision_transformer import PatchEmbed, Mlp
+from timm.layers import use_fused_attn
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+            self,
+            dim,
+            num_heads=8,
+            qkv_bias=False,
+            qk_norm=False,
+            attn_drop=0.,
+            proj_drop=0.,
+            norm_layer=nn.LayerNorm,
+            use_gamma=False
+    ):
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        if use_gamma:
+            self.gamma1 = nn.Parameter(torch.ones(dim * 3))
+            self.gamma2 = nn.Parameter(torch.ones(dim))
+        else:
+            self.gamma1 = 1
+            self.gamma2 = 1
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (self.gamma1 * self.qkv(x)).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if self.fused_attn:
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            )
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.gamma2 * self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size, dropout_prob):
+        super().__init__()
+        use_cfg_embedding = dropout_prob > 0
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+    def token_drop(self, labels, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        labels = torch.where(drop_ids, self.num_classes, labels)
+        return labels
+
+    def forward(self, labels, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            labels = self.token_drop(labels, force_drop_ids)
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+
+#################################################################################
+#                                 Core DiT Model                                #
+#################################################################################
+
+class DiTBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, use_gamma=False, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, use_gamma=use_gamma, **block_kwargs)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+        if use_gamma:
+            self.gamma1 = nn.Parameter(torch.ones(hidden_size))
+            self.gamma2 = nn.Parameter(torch.ones(hidden_size))
+        else:
+            self.gamma1 = 1
+            self.gamma2 = 1
+
+    def forward(self, x, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+        x = x + self.gamma1 * gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + self.gamma2 * gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class DiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        input_size=32,
+        patch_size=2,
+        in_channels=4,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        class_dropout_prob=0.1,
+        num_classes=1000,
+        learn_sigma=True,
+    ):
+        super().__init__()
+        self.learn_sigma = learn_sigma
+        self.in_channels = in_channels
+        self.out_channels = in_channels * 2 if learn_sigma else in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
+        num_patches = self.x_embedder.num_patches
+        # Will use fixed sin-cos embedding:
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+
+        use_gamma = [True] * 14 + [False] * 14
+        self.blocks = nn.ModuleList([
+            DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, use_gamma=use_gamma[depth_index]) for depth_index in range(depth)
+        ])
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5))
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        return imgs
+
+    def forward(self, x, t, y):
+        """
+        Forward pass of DiT.
+        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        y: (N,) tensor of class labels
+        """
+        x = self.x_embedder(x) + self.pos_embed  # (N, T, D), where T = H * W / patch_size ** 2
+        t = self.t_embedder(t)                   # (N, D)
+        y = self.y_embedder(y, self.training)    # (N, D)
+        c = t + y                                # (N, D)
+        for block in self.blocks:
+            x = block(x, c)                      # (N, T, D)
+        x = self.final_layer(x, c)                # (N, T, patch_size ** 2 * out_channels)
+        x = self.unpatchify(x)                   # (N, out_channels, H, W)
+        return x
+
+    def forward_with_cfg(self, x, t, y, cfg_scale, **kwargs):
+        """
+        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
+        half = x[: len(x) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self.forward(combined, t, y)
+        # For exact reproducibility reasons, we apply classifier-free guidance on only
+        # three channels by default. The standard approach to cfg applies it to all channels.
+        # This can be done by uncommenting the following line and commenting-out the line following that.
+        # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
+        eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+
+
+#################################################################################
+#                   Sine/Cosine Positional Embedding Functions                  #
+#################################################################################
+# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+#################################################################################
+#                                   DiT Configs                                  #
+#################################################################################
+
+def DiT_XL_2(**kwargs):
+    return DiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
+
+def DiT_XL_4(**kwargs):
+    return DiT(depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs)
+
+def DiT_XL_8(**kwargs):
+    return DiT(depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs)
+
+def DiT_L_2(**kwargs):
+    return DiT(depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs)
+
+def DiT_L_4(**kwargs):
+    return DiT(depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs)
+
+def DiT_L_8(**kwargs):
+    return DiT(depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs)
+
+def DiT_B_2(**kwargs):
+    return DiT(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs)
+
+def DiT_B_4(**kwargs):
+    return DiT(depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs)
+
+def DiT_B_8(**kwargs):
+    return DiT(depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs)
+
+def DiT_S_2(**kwargs):
+    return DiT(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs)
+
+def DiT_S_4(**kwargs):
+    return DiT(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs)
+
+def DiT_S_8(**kwargs):
+    return DiT(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs)
+
+
+DiT_models = {
+    'DiT-XL/2': DiT_XL_2,  'DiT-XL/4': DiT_XL_4,  'DiT-XL/8': DiT_XL_8,
+    'DiT-L/2':  DiT_L_2,   'DiT-L/4':  DiT_L_4,   'DiT-L/8':  DiT_L_8,
+    'DiT-B/2':  DiT_B_2,   'DiT-B/4':  DiT_B_4,   'DiT-B/8':  DiT_B_8,
+    'DiT-S/2':  DiT_S_2,   'DiT-S/4':  DiT_S_4,   'DiT-S/8':  DiT_S_8,
+}

var/D3HR/validation/models/mobilenet_v2.py

@@ -0,0 +1,151 @@
+"""
+Creates a MobileNetV2 Model as defined in:
+Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen. (2018). 
+MobileNetV2: Inverted Residuals and Linear Bottlenecks
+arXiv preprint arXiv:1801.04381.
+import from https://github.com/tonylins/pytorch-mobilenet-v2
+"""
+
+import torch.nn as nn
+import math
+
+__all__ = ['mobilenetv2']
+
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+def conv_3x3_bn(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.identity = stride == 1 and inp == oup
+
+        if expand_ratio == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.identity:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, num_classes=1000, width_mult=1.):
+        super(MobileNetV2, self).__init__()
+        # setting of inverted residual blocks
+        self.cfgs = [
+            # t, c, n, s
+            [1,  16, 1, 1],
+            [6,  24, 2, 2],
+            [6,  32, 3, 2],
+            [6,  64, 4, 2],
+            [6,  96, 3, 1],
+            [6, 160, 3, 2],
+            [6, 320, 1, 1],
+        ]
+
+        # building first layer
+        input_channel = _make_divisible(32 * width_mult, 4 if width_mult == 0.1 else 8)
+        layers = [conv_3x3_bn(3, input_channel, 2)]
+        # building inverted residual blocks
+        block = InvertedResidual
+        for t, c, n, s in self.cfgs:
+            output_channel = _make_divisible(c * width_mult, 4 if width_mult == 0.1 else 8)
+            for i in range(n):
+                layers.append(block(input_channel, output_channel, s if i == 0 else 1, t))
+                input_channel = output_channel
+        self.features = nn.Sequential(*layers)
+        # building last several layers
+        output_channel = _make_divisible(1280 * width_mult, 4 if width_mult == 0.1 else 8) if width_mult > 1.0 else 1280
+        self.conv = conv_1x1_bn(input_channel, output_channel)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.classifier = nn.Linear(output_channel, num_classes)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.conv(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+
+def mobilenetv2(**kwargs):
+    """
+    Constructs a MobileNet V2 model
+    """
+    return MobileNetV2(**kwargs)

var/D3HR/validation/models/pipeline_stable_unclip_img2img.py

@@ -0,0 +1,854 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import PIL.Image
+import torch
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.models.embeddings import get_timestep_embedding
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    deprecate,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import requests
+        >>> import torch
+        >>> from PIL import Image
+        >>> from io import BytesIO
+
+        >>> from diffusers import StableUnCLIPImg2ImgPipeline
+
+        >>> pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
+        ...     "fusing/stable-unclip-2-1-l-img2img", torch_dtype=torch.float16
+        ... )  # TODO update model path
+        >>> pipe = pipe.to("cuda")
+
+        >>> url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+
+        >>> response = requests.get(url)
+        >>> init_image = Image.open(BytesIO(response.content)).convert("RGB")
+        >>> init_image = init_image.resize((768, 512))
+
+        >>> prompt = "A fantasy landscape, trending on artstation"
+
+        >>> images = pipe(prompt, init_image).images
+        >>> images[0].save("fantasy_landscape.png")
+        ```
+"""
+
+
+class StableUnCLIPImg2ImgPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin):
+    """
+    Pipeline for text-guided image-to-image generation using stable unCLIP.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    The pipeline also inherits the following loading methods:
+        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
+        - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
+        - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
+
+    Args:
+        feature_extractor ([`CLIPImageProcessor`]):
+            Feature extractor for image pre-processing before being encoded.
+        image_encoder ([`CLIPVisionModelWithProjection`]):
+            CLIP vision model for encoding images.
+        image_normalizer ([`StableUnCLIPImageNormalizer`]):
+            Used to normalize the predicted image embeddings before the noise is applied and un-normalize the image
+            embeddings after the noise has been applied.
+        image_noising_scheduler ([`KarrasDiffusionSchedulers`]):
+            Noise schedule for adding noise to the predicted image embeddings. The amount of noise to add is determined
+            by the `noise_level`.
+        tokenizer (`~transformers.CLIPTokenizer`):
+            A [`~transformers.CLIPTokenizer`)].
+        text_encoder ([`~transformers.CLIPTextModel`]):
+            Frozen [`~transformers.CLIPTextModel`] text-encoder.
+        unet ([`UNet2DConditionModel`]):
+            A [`UNet2DConditionModel`] to denoise the encoded image latents.
+        scheduler ([`KarrasDiffusionSchedulers`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+    """
+
+    model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
+    _exclude_from_cpu_offload = ["image_normalizer"]
+
+    # image encoding components
+    feature_extractor: CLIPImageProcessor
+    image_encoder: CLIPVisionModelWithProjection
+
+    # image noising components
+    image_normalizer: StableUnCLIPImageNormalizer
+    image_noising_scheduler: KarrasDiffusionSchedulers
+
+    # regular denoising components
+    tokenizer: CLIPTokenizer
+    text_encoder: CLIPTextModel
+    unet: UNet2DConditionModel
+    scheduler: KarrasDiffusionSchedulers
+
+    vae: AutoencoderKL
+
+    def __init__(
+        self,
+        # image encoding components
+        feature_extractor: CLIPImageProcessor,
+        image_encoder: CLIPVisionModelWithProjection,
+        # image noising components
+        image_normalizer: StableUnCLIPImageNormalizer,
+        image_noising_scheduler: KarrasDiffusionSchedulers,
+        # regular denoising components
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModel,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        # vae
+        vae: AutoencoderKL,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            feature_extractor=feature_extractor,
+            image_encoder=image_encoder,
+            image_normalizer=image_normalizer,
+            image_noising_scheduler=image_noising_scheduler,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            unet=unet,
+            scheduler=scheduler,
+            vae=vae,
+        )
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+        **kwargs,
+    ):
+        deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
+        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
+
+        prompt_embeds_tuple = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=lora_scale,
+            **kwargs,
+        )
+
+        # concatenate for backwards comp
+        prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
+
+        return prompt_embeds
+
+    def _encode_image(
+        self,
+        image,
+        device,
+        batch_size,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        noise_level,
+        generator,
+        image_embeds,
+    ):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if isinstance(image, PIL.Image.Image):
+            # the image embedding should repeated so it matches the total batch size of the prompt
+            repeat_by = batch_size
+        else:
+            # assume the image input is already properly batched and just needs to be repeated so
+            # it matches the num_images_per_prompt.
+            #
+            # NOTE(will) this is probably missing a few number of side cases. I.e. batched/non-batched
+            # `image_embeds`. If those happen to be common use cases, let's think harder about
+            # what the expected dimensions of inputs should be and how we handle the encoding.
+            repeat_by = num_images_per_prompt
+
+        if image_embeds is None:
+            if not isinstance(image, torch.Tensor):
+                image = self.feature_extractor(images=image, return_tensors="pt").pixel_values
+
+            image = image.to(device=device, dtype=dtype)
+            image_embeds = self.image_encoder(image).image_embeds
+
+        image_embeds = self.noise_image_embeddings(
+            image_embeds=image_embeds,
+            noise_level=noise_level,
+            generator=generator,
+        )
+
+        # duplicate image embeddings for each generation per prompt, using mps friendly method
+        image_embeds = image_embeds.unsqueeze(1)
+        bs_embed, seq_len, _ = image_embeds.shape
+        image_embeds = image_embeds.repeat(1, repeat_by, 1)
+        image_embeds = image_embeds.view(bs_embed * repeat_by, seq_len, -1)
+        image_embeds = image_embeds.squeeze(1)
+
+        if do_classifier_free_guidance:
+            negative_prompt_embeds = torch.zeros_like(image_embeds)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            image_embeds = torch.cat([negative_prompt_embeds, image_embeds])
+
+        return image_embeds
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+        clip_skip: Optional[int] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            lora_scale (`float`, *optional*):
+                A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+        """
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if not USE_PEFT_BACKEND:
+                adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
+            else:
+                scale_lora_layers(self.text_encoder, lora_scale)
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+                )
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            if clip_skip is None:
+                prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
+                prompt_embeds = prompt_embeds[0]
+            else:
+                prompt_embeds = self.text_encoder(
+                    text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
+                )
+                # Access the `hidden_states` first, that contains a tuple of
+                # all the hidden states from the encoder layers. Then index into
+                # the tuple to access the hidden states from the desired layer.
+                prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]
+                # We also need to apply the final LayerNorm here to not mess with the
+                # representations. The `last_hidden_states` that we typically use for
+                # obtaining the final prompt representations passes through the LayerNorm
+                # layer.
+                prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)
+
+        if self.text_encoder is not None:
+            prompt_embeds_dtype = self.text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:
+            # Retrieve the original scale by scaling back the LoRA layers
+            unscale_lora_layers(self.text_encoder, lora_scale)
+
+        return prompt_embeds, negative_prompt_embeds
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
+    def decode_latents(self, latents):
+        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
+        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents, return_dict=False)[0]
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt,
+        image,
+        height,
+        width,
+        callback_steps,
+        noise_level,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        image_embeds=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Please make sure to define only one of the two."
+            )
+
+        if prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+
+        if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                "Provide either `negative_prompt` or `negative_prompt_embeds`. Cannot leave both `negative_prompt` and `negative_prompt_embeds` undefined."
+            )
+
+        if prompt is not None and negative_prompt is not None:
+            if type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        if noise_level < 0 or noise_level >= self.image_noising_scheduler.config.num_train_timesteps:
+            raise ValueError(
+                f"`noise_level` must be between 0 and {self.image_noising_scheduler.config.num_train_timesteps - 1}, inclusive."
+            )
+
+        if image is not None and image_embeds is not None:
+            raise ValueError(
+                "Provide either `image` or `image_embeds`. Please make sure to define only one of the two."
+            )
+
+        if image is None and image_embeds is None:
+            raise ValueError(
+                "Provide either `image` or `image_embeds`. Cannot leave both `image` and `image_embeds` undefined."
+            )
+
+        if image is not None:
+            if (
+                not isinstance(image, torch.Tensor)
+                and not isinstance(image, PIL.Image.Image)
+                and not isinstance(image, list)
+            ):
+                raise ValueError(
+                    "`image` has to be of type `torch.FloatTensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+                    f" {type(image)}"
+                )
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_unclip.StableUnCLIPPipeline.noise_image_embeddings
+    def noise_image_embeddings(
+        self,
+        image_embeds: torch.Tensor,
+        noise_level: int,
+        noise: Optional[torch.FloatTensor] = None,
+        generator: Optional[torch.Generator] = None,
+    ):
+        """
+        Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher
+        `noise_level` increases the variance in the final un-noised images.
+
+        The noise is applied in two ways:
+        1. A noise schedule is applied directly to the embeddings.
+        2. A vector of sinusoidal time embeddings are appended to the output.
+
+        In both cases, the amount of noise is controlled by the same `noise_level`.
+
+        The embeddings are normalized before the noise is applied and un-normalized after the noise is applied.
+        """
+        if noise is None:
+            noise = randn_tensor(
+                image_embeds.shape, generator=generator, device=image_embeds.device, dtype=image_embeds.dtype
+            )
+
+        noise_level = torch.tensor([noise_level] * image_embeds.shape[0], device=image_embeds.device)
+
+        self.image_normalizer.to(image_embeds.device)
+        image_embeds = self.image_normalizer.scale(image_embeds)
+
+        image_embeds = self.image_noising_scheduler.add_noise(image_embeds, timesteps=noise_level, noise=noise)
+
+        image_embeds = self.image_normalizer.unscale(image_embeds)
+
+        noise_level = get_timestep_embedding(
+            timesteps=noise_level, embedding_dim=image_embeds.shape[-1], flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+
+        # `get_timestep_embeddings` does not contain any weights and will always return f32 tensors,
+        # but we might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        noise_level = noise_level.to(image_embeds.dtype)
+
+        image_embeds = torch.cat((image_embeds, noise_level), 1)
+
+        return image_embeds
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: Union[torch.FloatTensor, PIL.Image.Image] = None,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 20,
+        guidance_scale: float = 10,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        noise_level: int = 0,
+        image_embeds: Optional[torch.FloatTensor] = None,
+        clip_skip: Optional[int] = None,
+        cond_fn = None,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, either `prompt_embeds` will be
+                used or prompt is initialized to `""`.
+            image (`torch.FloatTensor` or `PIL.Image.Image`):
+                `Image` or tensor representing an image batch. The image is encoded to its CLIP embedding which the
+                `unet` is conditioned on. The image is _not_ encoded by the `vae` and then used as the latents in the
+                denoising process like it is in the standard Stable Diffusion text-guided image variation process.
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 20):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 10.0):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            noise_level (`int`, *optional*, defaults to `0`):
+                The amount of noise to add to the image embeddings. A higher `noise_level` increases the variance in
+                the final un-noised images. See [`StableUnCLIPPipeline.noise_image_embeddings`] for more details.
+            image_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated CLIP embeddings to condition the `unet` on. These latents are not used in the denoising
+                process. If you want to provide pre-generated latents, pass them to `__call__` as `latents`.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`:
+                [`~ pipeline_utils.ImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When returning
+                a tuple, the first element is a list with the generated images.
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        if prompt is None and prompt_embeds is None:
+            prompt = len(image) * [""] if isinstance(image, list) else ""
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt=prompt,
+            image=image,
+            height=height,
+            width=width,
+            callback_steps=callback_steps,
+            noise_level=noise_level,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            image_embeds=image_embeds,
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        batch_size = batch_size * num_images_per_prompt
+
+        device = self._execution_device
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        # 4. Encoder input image
+        noise_level = torch.tensor([noise_level], device=device)
+        image_embeds = self._encode_image(
+            image=image,
+            device=device,
+            batch_size=batch_size,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            noise_level=noise_level,
+            generator=generator,
+            image_embeds=image_embeds,
+        )
+
+        # 5. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 6. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size=batch_size,
+            num_channels_latents=num_channels_latents,
+            height=height,
+            width=width,
+            dtype=prompt_embeds.dtype,
+            device=device,
+            generator=generator,
+            latents=latents,
+        )
+
+        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 8. Denoising loop
+        for i, t in enumerate(timesteps):
+            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            # predict the noise residual
+            noise_pred = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=prompt_embeds,
+                class_labels=image_embeds,
+                cross_attention_kwargs=cross_attention_kwargs,
+                return_dict=False,
+            )[0]
+
+            # perform guidance
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents, cond_fn, **extra_step_kwargs, return_dict=False)[0]
+
+            if callback is not None and i % callback_steps == 0:
+                step_idx = i // getattr(self.scheduler, "order", 1)
+                callback(step_idx, t, latents)
+
+        # 9. Post-processing
+        if not output_type == "latent":
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+        else:
+            image = latents
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)

var/D3HR/validation/models/resnet.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+from torchvision.models.resnet import (
+    ResNet, ResNet18_Weights, ResNet50_Weights, ResNet101_Weights,
+    BasicBlock, Bottleneck,
+    _ovewrite_named_param
+)
+
+
+class FeatResNet(ResNet):
+    def __init__(self, block, layers, **kwargs):
+        super(FeatResNet, self).__init__(block, layers, **kwargs)
+    
+    def get_features(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+
+        return x
+
+
+def resnet18(*, weights=None, progress=True, **kwargs):
+    weights = ResNet18_Weights.verify(weights)
+    if weights is not None:
+        _ovewrite_named_param(kwargs, 'num_classes', len(weights.meta['categories']))
+
+    model = FeatResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.get_state_dict(progress=progress))
+    
+    return model
+
+
+def resnet50(*, weights=None, progress=True, **kwargs):
+    weights = ResNet50_Weights.verify(weights)
+    if weights is not None:
+        _ovewrite_named_param(kwargs, 'num_classes', len(weights.meta['categories']))
+
+    model = FeatResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.get_state_dict(progress=progress))
+    
+    return model
+
+
+def resnet101(*, weights=None, progress=True, **kwargs):
+    weights = ResNet101_Weights.verify(weights)
+    if weights is not None:
+        _ovewrite_named_param(kwargs, 'num_classes', len(weights.meta['categories']))
+
+    model = FeatResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.get_state_dict(progress=progress))
+    
+    return model
+
+
+def resnet152(*, weights=None, progress=True, **kwargs):
+    weights = ResNet101_Weights.verify(weights)
+    if weights is not None:
+        _ovewrite_named_param(kwargs, 'num_classes', len(weights.meta['categories']))
+
+    model = FeatResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.get_state_dict(progress=progress))
+    
+    return model

var/D3HR/validation/models/scheduling_ddim.py

@@ -0,0 +1,522 @@
+# Copyright 2024 Stanford University Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+
+
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
+class DDIMSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.Tensor
+    pred_original_sample: Optional[torch.Tensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+
+
+    Args:
+        betas (`torch.Tensor`):
+            the betas that the scheduler is being initialized with.
+
+    Returns:
+        `torch.Tensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+
+class DDIMScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `DDIMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
+    non-Markovian guidance.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        clip_sample (`bool`, defaults to `True`):
+            Clip the predicted sample for numerical stability.
+        clip_sample_range (`float`, defaults to 1.0):
+            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
+        set_alpha_to_one (`bool`, defaults to `True`):
+            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
+            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
+            otherwise it uses the alpha value at step 0.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
+        timestep_spacing (`str`, defaults to `"leading"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        clip_sample: bool = True,
+        set_alpha_to_one: bool = True,
+        steps_offset: int = 0,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        clip_sample_range: float = 1.0,
+        sample_max_value: float = 1.0,
+        timestep_spacing: str = "leading",
+        rescale_betas_zero_snr: bool = False,
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+
+        # Rescale for zero SNR
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        # At every step in ddim, we are looking into the previous alphas_cumprod
+        # For the final step, there is no previous alphas_cumprod because we are already at 0
+        # `set_alpha_to_one` decides whether we set this parameter simply to one or
+        # whether we use the final alpha of the "non-previous" one.
+        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # setable values
+        self.num_inference_steps = None
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    def _get_variance(self, timestep, prev_timestep):
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
+
+        return variance
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+        """
+
+        if num_inference_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
+                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
+                f" maximal {self.config.num_train_timesteps} timesteps."
+            )
+
+        self.num_inference_steps = num_inference_steps
+
+        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
+        if self.config.timestep_spacing == "linspace":
+            timesteps = (
+                np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps)
+                .round()[::-1]
+                .copy()
+                .astype(np.int64)
+            )
+        elif self.config.timestep_spacing == "leading":
+            step_ratio = self.config.num_train_timesteps // self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+            timesteps += self.config.steps_offset
+        elif self.config.timestep_spacing == "trailing":
+            step_ratio = self.config.num_train_timesteps / self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'."
+            )
+
+        self.timesteps = torch.from_numpy(timesteps).to(device)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int,
+        sample: torch.Tensor,
+        cond_fn = None,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[DDIMSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            eta (`float`):
+                The weight of noise for added noise in diffusion step.
+            use_clipped_model_output (`bool`, defaults to `False`):
+                If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary
+                because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no
+                clipping has happened, "corrected" `model_output` would coincide with the one provided as input and
+                `use_clipped_model_output` has no effect.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            variance_noise (`torch.Tensor`):
+                Alternative to generating noise with `generator` by directly providing the noise for the variance
+                itself. Useful for methods such as [`CycleDiffusion`].
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
+        # Ideally, read DDIM paper in-detail understanding
+
+        # Notation (<variable name> -> <name in paper>
+        # - pred_noise_t -> e_theta(x_t, t)
+        # - pred_original_sample -> f_theta(x_t, t) or x_0
+        # - std_dev_t -> sigma_t
+        # - eta -> η
+        # - pred_sample_direction -> "direction pointing to x_t"
+        # - pred_prev_sample -> "x_t-1"
+
+        # 1. get previous step value (=t-1)
+        prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps
+
+        # 2. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+
+        beta_prod_t = 1 - alpha_prod_t
+
+        if cond_fn is not None:
+            model_output = model_output - (1 - alpha_prod_t) ** (0.5) * cond_fn(sample)
+
+        # 3. compute predicted original sample from predicted noise also called
+        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        if self.config.prediction_type == "epsilon":
+            pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+            pred_epsilon = model_output
+        elif self.config.prediction_type == "sample":
+            pred_original_sample = model_output
+            pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
+        elif self.config.prediction_type == "v_prediction":
+            pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
+            pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
+                " `v_prediction`"
+            )
+
+        # 4. Clip or threshold "predicted x_0"
+        if self.config.thresholding:
+            pred_original_sample = self._threshold_sample(pred_original_sample)
+        elif self.config.clip_sample:
+            pred_original_sample = pred_original_sample.clamp(
+                -self.config.clip_sample_range, self.config.clip_sample_range
+            )
+
+        # 5. compute variance: "sigma_t(η)" -> see formula (16)
+        # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
+        variance = self._get_variance(timestep, prev_timestep)
+        std_dev_t = eta * variance ** (0.5)
+
+        if use_clipped_model_output:
+            # the pred_epsilon is always re-derived from the clipped x_0 in Glide
+            pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
+
+        # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon
+
+        # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
+        prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
+
+        if eta > 0:
+            if variance_noise is not None and generator is not None:
+                raise ValueError(
+                    "Cannot pass both generator and variance_noise. Please make sure that either `generator` or"
+                    " `variance_noise` stays `None`."
+                )
+
+            if variance_noise is None:
+                variance_noise = randn_tensor(
+                    model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype
+                )
+            variance = std_dev_t * variance_noise
+
+            prev_sample = prev_sample + variance
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement
+        # for the subsequent add_noise calls
+        self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device)
+        alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
+    def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device)
+        alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype)
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(sample.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+        return velocity
+
+    def __len__(self):
+        return self.config.num_train_timesteps

var/D3HR/validation/utils/data_utils.py

@@ -0,0 +1,431 @@
+import os
+import random
+import numpy as np
+from PIL import Image
+import torch
+import torchvision
+from torchvision import transforms
+import json
+from torch.utils.data import Dataset
+
+    
+def find_subclasses(spec, nclass, phase=0):
+    classes = []
+    cls_from = nclass * phase
+    cls_to = nclass * (phase + 1)
+    if spec == 'woof':
+        file_list = './misc/class_woof.txt'
+    elif spec == 'im100':
+        file_list = './misc/class_100.txt'
+    else:
+        file_list = './misc/class_indices.txt'
+    with open(file_list, 'r') as f:
+        class_name = f.readlines()
+    for c in class_name:
+        c = c.split('\n')[0]
+        classes.append(c)
+    classes = classes[cls_from:cls_to]
+    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
+
+    return classes, class_to_idx
+
+
+def find_original_classes(spec, classes):
+    file_list = './misc/class_indices.txt'
+    with open(file_list, 'r') as f:
+        all_classes = f.readlines()
+    all_classes = [class_name.split('\n')[0] for class_name in all_classes]
+    original_classes = []
+    for class_name in classes:
+        original_classes.append(all_classes.index(class_name))
+    return original_classes
+
+
+def load_mapping_imgwoof(mapping_file, names):
+    new_mapping = {}
+    with open(mapping_file, 'r') as file:
+        data = json.load(file)
+        if "tiny" in mapping_file:
+            for index, line in enumerate(file):
+                # 提取每一行的编号（n开头部分）并将行号-1
+                key = line.split()[0]
+                new_mapping[key] = index 
+        else:
+            new_mapping = {item["wnid"]: names.index(item["name"]) for item in data.values() if item['name'] in names}
+    return new_mapping
+
+
+def load_mapping(mapping_file):
+    new_mapping = {}
+    with open(mapping_file, 'r') as file:
+        data = json.load(file)
+        if "tiny" in mapping_file:
+            for index, line in enumerate(file):
+                # 提取每一行的编号（n开头部分）并将行号-1
+                key = line.split()[0]
+                new_mapping[key] = index 
+        else:
+            new_mapping = {item["wnid"]: item["index"] for item in data.values()}
+    return new_mapping
+
+
+
+def load_mapping_txt(mapping_file):
+        wnid_to_index = {}
+        with open(mapping_file, 'r') as f:
+            for line in f:
+                wnid, index = line.strip().split('\t')
+                wnid_to_index[wnid] = int(index)
+        return wnid_to_index
+
+def find_classes(class_file):
+    with open(class_file) as r:
+        classes = list(map(lambda s: s.strip(), r.readlines()))
+
+    classes.sort()
+    class_to_idx = {classes[i]: i for i in range(len(classes))}
+
+    return class_to_idx
+
+class ImageFolder(Dataset):
+    def __init__(self, split=None, txt_file=None, subset=None, mapping_file=None, transform=None):
+        super(ImageFolder, self).__init__()
+        self.split = split
+        self.image_paths = []
+        self.targets = []
+        self.samples = []
+        self.subset = subset
+        if self.subset == 'imagenet_1k':
+            self.wnid_to_index = load_mapping(mapping_file)
+        elif self.subset == 'tinyimagenet':
+            self.wnid_to_index = find_classes(mapping_file)
+        if split == 'train':
+            self._load_from_txt(txt_file)
+        else:
+            self._load_from_txt(txt_file)
+        self.transform = transform
+
+
+    def _load_from_txt(self, txt_file):
+        with open(txt_file, "r") as file:
+            image_paths = file.readlines()
+
+        # 去掉每行的换行符
+        self.image_paths = [path.strip() for path in image_paths]
+        for path in self.image_paths:
+            self.samples.append(path)
+            if self.subset == 'cifar10' or self.subset == 'cifar100':
+                class_index = int(path.split('/')[-2][-3:])
+            else:
+                # if self.split == 'test':
+                #     class_index = self.wnid_to_index[path.split('/')[-2]]
+                # elif self.split == 'train':
+                class_index = self.wnid_to_index[path.split('/')[-2]]
+            self.targets.append(class_index)
+
+
+    #combine ten txt
+    def _load_from_txt_1(self, txt_file):
+
+        image_paths_10 = []
+        for kk in range(10):
+            txt_file=f'/scratch/zhao.lin1/tinyimagenet_finetune_start_step_18_ddim_inversion_10_min_images_{kk}/train.txt'
+            with open(txt_file, "r") as file:
+                image_paths = file.readlines()
+
+            image_paths_10.append([path.strip() for path in image_paths])
+
+
+        for kk in range(len(image_paths)):
+            number = random.randint(0, 9)
+            self.image_paths.append(image_paths_10[number][kk])
+            if self.subset == 'cifar10' or self.subset == 'cifar100':
+                class_index = int(path.split('/')[-2][-3:])
+            else:
+                # if self.split == 'test':
+                #     class_index = self.wnid_to_index[path.split('/')[-2]]
+                # elif self.split == 'train':
+                class_index = self.wnid_to_index[image_paths_10[number][kk].split('/')[-2]]
+            self.targets.append(class_index)
+
+    def __getitem__(self, index):
+        img_path = self.image_paths[index]
+        try:
+            sample = Image.open(img_path).convert('RGB')
+        except Exception as e:
+            print(f"Error loading image {img_path}: {e}")
+            # Return a black image in case of error
+            sample = Image.new('RGB', (256, 256))
+        sample = self.transform(sample)
+        # class_dir = img_path.split('/')[-2]
+        return sample, self.targets[index]
+
+    def __len__(self):
+        return len(self.targets)
+
+
+class Imagewoof(Dataset):
+    def __init__(self, split=None, txt_file=None, subset=None, mapping_file=None, transform=None):
+        super(Imagewoof, self).__init__()
+        self.split = split
+        self.image_paths = []
+        self.targets = []
+        self.samples = []
+        self.subset = subset
+        self.names = ["Australian_terrier", "Border_terrier", "Samoyed", "beagle", "Shih-Tzu", "English_foxhound", "Rhodesian_ridgeback", "dingo", "golden_retriever", "Old_English_sheepdog"]
+        self.wnid_to_index = load_mapping_imgwoof(mapping_file, self.names)
+        self._load_from_txt(txt_file)
+        self.transform = transform
+
+
+    def _load_from_txt(self, txt_file):
+        with open(txt_file, "r") as file:
+            image_paths = file.readlines()
+
+        # 去掉每行的换行符
+        image_paths = [path.strip() for path in image_paths]
+        for path in image_paths:
+            self.samples.append(path)
+            if self.subset == 'cifar10' or self.subset == 'cifar100':
+                class_index = int(path.split('/')[-2][-3:])
+            else:
+                # if self.split == 'test':
+                #     class_index = self.wnid_to_index[path.split('/')[-2]]
+                # elif self.split == 'train':
+                if path.split('/')[-2] in list(self.wnid_to_index.keys()):
+                    class_index = self.wnid_to_index[path.split('/')[-2]]
+                    self.image_paths.append(path)
+                    self.targets.append(class_index)
+
+    def __getitem__(self, index):
+        img_path = self.image_paths[index]
+        try:
+            sample = Image.open(img_path).convert('RGB')
+        except Exception as e:
+            print(f"Error loading image {img_path}: {e}")
+            # Return a black image in case of error
+            sample = Image.new('RGB', (256, 256))
+        sample = self.transform(sample)
+        # class_dir = img_path.split('/')[-2]
+        return sample, self.targets[index]
+
+    def __len__(self):
+        return len(self.targets)
+
+
+
+# class ImageFolder(torchvision.datasets.ImageFolder):
+#     def __init__(self, nclass, ipc, mem=False, spec='none', phase=0, **kwargs):
+#         super(ImageFolder, self).__init__(**kwargs)
+#         self.mem = mem
+#         self.spec = spec
+#         self.classes, self.class_to_idx = find_subclasses(
+#             spec=spec, nclass=nclass, phase=phase
+#         )
+#         self.original_classes = find_original_classes(spec=self.spec, classes=self.classes)
+#         self.samples, self.targets = self.load_subset(ipc=ipc)
+#         if self.mem:
+#             self.samples = [self.loader(path) for path in self.samples]
+
+#     def load_subset(self, ipc=-1):
+#         all_samples = torchvision.datasets.folder.make_dataset(
+#             self.root, self.class_to_idx, self.extensions
+#         )
+#         samples = np.array([item[0] for item in all_samples])
+#         targets = np.array([item[1] for item in all_samples])
+
+#         if ipc == -1:
+#             return samples, targets
+#         else:
+#             sub_samples = []
+#             sub_targets = []
+#             for c in range(len(self.classes)):
+#                 c_indices = np.where(targets == c)[0]
+#                 #random.shuffle(c_indices)
+#                 sub_samples.extend(samples[c_indices[:ipc]])
+#                 sub_targets.extend(targets[c_indices[:ipc]])
+#             return sub_samples, sub_targets
+
+#     def __getitem__(self, index):
+#         if self.mem:
+#             sample = self.samples[index]
+#         else:
+#             sample = self.loader(self.samples[index])
+#         sample = self.transform(sample)
+#         return sample, self.targets[index]
+
+#     def __len__(self):
+#         return len(self.targets)
+def random_stitch_crop_4(image):
+    """随机从 stitch 的四个子区域中裁剪一个"""
+    w, h = image.size  # 获取图像的宽和高
+    w_half, h_half = w // 2, h // 2
+    
+    # 定义四个区域的坐标
+    regions = [
+        (0, 0, w_half, h_half),  # 左上
+        (w_half, 0, w, h_half),  # 右上
+        (0, h_half, w_half, h),  # 左下
+        (w_half, h_half, w, h),  # 右下
+    ]
+    
+    # 随机选择一个区域
+    x1, y1, x2, y2 = random.choice(regions)
+    return image.crop((x1, y1, x2, y2))  # 裁��并返回
+
+def transform_imagenet(args):
+    resize_test = [transforms.Resize(args.input_size // 7 * 8), transforms.CenterCrop(args.input_size)]
+    # resize_test = [transforms.Resize(args.input_size), transforms.CenterCrop(args.input_size)]
+
+    cast = [transforms.ToTensor()]
+
+    aug = [
+        # transforms.Resize(224), 
+        # transforms.Lambda(random_stitch_crop_4),
+        # ShufflePatches(args.factor),
+        transforms.RandomResizedCrop(
+            size=args.input_size,
+            # scale=(0.5, 1.0),
+            # scale=(1 / args.factor, args.max_scale_crops),
+            scale=(0.08, args.max_scale_crops),
+            antialias=True,
+        ),
+        transforms.RandomHorizontalFlip()
+    ]
+
+    normalize = [transforms.Normalize(
+        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+    )]
+
+    train_transform = transforms.Compose(aug +cast+ normalize)
+    test_transform = transforms.Compose(resize_test + cast + normalize)
+
+    return train_transform, test_transform
+
+
+sharing_strategy = "file_system"
+torch.multiprocessing.set_sharing_strategy(sharing_strategy)
+
+
+def set_worker_sharing_strategy(worker_id: int) -> None:
+    torch.multiprocessing.set_sharing_strategy(sharing_strategy)
+
+
+def load_data(args, coreset=False, resize_only=False, mem_flag=True, trainset_only=False):
+    train_transform, test_transform = transform_imagenet(args)
+    # if len(args.data_dir) == 1:
+    #     train_dir = os.path.join(args.data_dir[0], 'train')
+    #     val_dir = os.path.join(args.data_dir[0], 'val')
+    # else:
+    #     train_dir = args.data_dir[0]
+    #     val_dir = os.path.join(args.data_dir[1], 'val')
+
+    if resize_only:
+        train_transform = transforms.Compose([
+            transforms.Resize((512, 512)),
+        ])
+    elif coreset:
+        train_transform = test_transform
+
+    # train_dataset = ImageFolder(
+    #     nclass=args.nclass,
+    #     ipc=args.ipc,
+    #     mem=mem_flag,
+    #     spec=args.spec,
+    #     phase=args.phase,
+    #     root=train_dir,
+    #     transform=train_transform,
+    # )
+
+    if args.subset == 'imagewoof':
+        # Imagewoor
+        train_dataset = Imagewoof(
+            split = 'train',
+            txt_file=args.txt_file,
+            mapping_file=args.mapping_file,
+            subset = args.subset,
+            transform=train_transform,
+        )  
+    else:
+        train_dataset = ImageFolder(
+            split = 'train',
+            txt_file=args.txt_file,
+            mapping_file=args.mapping_file,
+            subset = args.subset,
+            transform=train_transform,
+        )
+
+
+
+    if trainset_only:
+        return train_dataset
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=24,
+        pin_memory=True,
+        worker_init_fn=set_worker_sharing_strategy,
+    )
+    if args.subset == 'cifar10':
+        val_dataset = torchvision.datasets.CIFAR10(root='/scratch/zhao.lin1/', train=False, download=True, transform=test_transform)
+    elif args.subset == 'cifar100':
+        val_dataset = torchvision.datasets.CIFAR100(root='/scratch/zhao.lin1/', train=False, download=True, transform=test_transform)
+    elif args.subset == 'imagewoof':
+        val_dataset = Imagewoof(
+                split = 'test',
+                txt_file=args.val_txt_file,
+                mapping_file=args.mapping_file,
+                subset = args.subset,
+                transform=test_transform,
+            )
+    else:
+        val_dataset = ImageFolder(
+                split = 'test',
+                txt_file=args.val_txt_file,
+                mapping_file=args.mapping_file,
+                subset = args.subset,
+                transform=test_transform,
+            )
+
+
+
+    val_loader = torch.utils.data.DataLoader(
+        val_dataset,
+        batch_size=256,
+        shuffle=False,
+        num_workers=24,
+        pin_memory=True,
+        worker_init_fn=set_worker_sharing_strategy,
+    )
+    print("load data successfully")
+
+    return train_dataset, train_loader, val_loader
+
+
+class ShufflePatches(torch.nn.Module):
+    def __init__(self, factor):
+        super().__init__()
+        self.factor = factor
+
+    def shuffle_weight(self, img, factor):
+        h, w = img.shape[1:]
+        tw = w // factor
+        patches = []
+        for i in range(factor):
+            i = i * tw
+            if i != factor - 1:
+                patches.append(img[..., i : i + tw])
+            else:
+                patches.append(img[..., i:])
+        random.shuffle(patches)
+        img = torch.cat(patches, -1)
+        return img
+
+    def forward(self, img):
+        img = self.shuffle_weight(img, self.factor)
+        img = img.permute(0, 2, 1)
+        img = self.shuffle_weight(img, self.factor)
+        img = img.permute(0, 2, 1)
+        return img

var/D3HR/validation/utils/download.py

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Functions for downloading pre-trained DiT models
+"""
+from torchvision.datasets.utils import download_url
+import torch
+import os
+
+
+pretrained_models = {'DiT-XL-2-512x512.pt', 'DiT-XL-2-256x256.pt'}
+
+
+def find_model(model_name):
+    """
+    Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path.
+    """
+    if model_name in pretrained_models:  # Find/download our pre-trained DiT checkpoints
+        return download_model(model_name)
+    else:  # Load a custom DiT checkpoint:
+        assert os.path.isfile(model_name), f'Could not find DiT checkpoint at {model_name}'
+        checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
+        if "ema" in checkpoint:  # supports checkpoints from train.py
+            checkpoint = checkpoint["ema"]
+        return checkpoint
+
+
+def download_model(model_name):
+    """
+    Downloads a pre-trained DiT model from the web.
+    """
+    assert model_name in pretrained_models
+    local_path = f'pretrained_models/{model_name}'
+    if not os.path.isfile(local_path):
+        os.makedirs('pretrained_models', exist_ok=True)
+        web_path = f'https://dl.fbaipublicfiles.com/DiT/models/{model_name}'
+        download_url(web_path, 'pretrained_models')
+    model = torch.load(local_path, map_location=lambda storage, loc: storage)
+    return model
+
+
+if __name__ == "__main__":
+    # Download all DiT checkpoints
+    for model in pretrained_models:
+        download_model(model)
+    print('Done.')

var/D3HR/validation/utils/syn_utils_dit.py

@@ -0,0 +1,172 @@
+import json
+import random
+import numpy as np
+
+import torch
+import torch.nn as nn
+from torchvision import transforms as T
+from torchvision.models import resnet18
+from transformers import CLIPModel, AutoTokenizer
+
+from .download import find_model
+from diffusion import create_diffusion
+from models.dit_models import DiT_models
+from diffusers.models import AutoencoderKL
+
+
+class SupConLoss(nn.Module):
+    def __init__(self, temperature=0.05, base_temperatue=0.05):
+        super(SupConLoss, self).__init__()
+        self.temperature = temperature
+        self.base_temperature = base_temperatue
+    
+    def forward(self, image_features, text_features, text_labels):
+        logits = (image_features @ text_features.T) / self.temperature
+        logits_max, _ = torch.max(logits, dim=1, keepdim=True)
+        logits = logits - logits_max.detach()
+
+        exp_logits = torch.exp(logits) * text_labels
+        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
+        mean_log_prob_pos = ((1 - text_labels) * log_prob).sum(1) / (1 - text_labels).sum(1)
+        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+        loss = loss.mean()
+
+        return loss
+
+
+class ImageSynthesizer(object):
+    def __init__(self, args):
+        self.vae = AutoencoderKL.from_pretrained(args.vae_path).to('cuda')
+        self.clip_model = CLIPModel.from_pretrained('laion/CLIP-ViT-L-14-laion2B-s32B-b82K').to('cuda')
+        self.clip_tokenizer = AutoTokenizer.from_pretrained('laion/CLIP-ViT-L-14-laion2B-s32B-b82K')
+
+        # DiT model
+        assert args.dit_image_size % 8 == 0, 'Image size must be divisible by 8'
+        latent_size = args.dit_image_size // 8
+        self.latent_size = latent_size
+        self.dit = DiT_models[args.dit_model](
+            input_size=latent_size,
+            num_classes=args.num_dit_classes
+        ).to('cuda')
+        ckpt_path = args.ckpt
+        state_dict = find_model(ckpt_path)
+        self.dit.load_state_dict(state_dict, strict=False)
+
+        # Diffusion
+        self.diffusion = create_diffusion(str(args.diffusion_steps))
+
+        # Class description
+        self.description_file = args.description_path
+        self.load_class_description()
+
+        self.cfg_scale = args.cfg_scale
+        self.clip_alpha = args.clip_alpha
+        self.cls_alpha = args.cls_alpha
+        self.num_pos_samples = 5
+        self.num_neg_samples = args.num_neg_samples
+        self.clip_normalize = T.Normalize(
+            mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]
+        )
+        self.contrast_criterion = SupConLoss()
+        self.neg_policy = args.neg_policy
+
+    def load_class_description(self):
+        with open(self.description_file, 'r') as fp:
+            descriptions = json.load(fp)
+        self.class_names = {}
+        self.descriptions = {}
+
+        for class_index, (class_name, description) in descriptions.items():
+            self.class_names[class_index] = class_name
+            self.descriptions[class_index] = description
+        
+        self.class_indices = list(self.class_names.keys())
+        self.class_name_list = list(self.class_names.values())
+        self.neighbors = {}
+        with torch.no_grad():
+            class_name_feat = self.extract_text_feature(self.class_name_list)
+            name_sims = (class_name_feat @ class_name_feat.T).cpu()
+            name_sims -= torch.eye(len(name_sims))
+            name_sims = name_sims.numpy()
+        for class_index, sim_indices in zip(self.class_indices, name_sims):
+            self.neighbors[class_index] = list(sim_indices)
+
+    def extract_text_feature(self, descriptions):
+        input_text = self.clip_tokenizer(descriptions, padding=True, return_tensors='pt').to('cuda')
+        text_feature = self.clip_model.get_text_features(**input_text)
+        text_feature = text_feature / text_feature.norm(dim=-1, keepdim=True)
+        return text_feature
+
+    def cond_fn(self, x, t, y=None, text_features=None, contrastive=True, **kwargs):
+        with torch.enable_grad():
+            x = nn.Parameter(x).requires_grad_()
+            real_x, _ = x.chunk(2, dim=0)
+            pseudo_image = self.vae.decode(real_x / 0.18215, return_dict=False)[0]
+            pseudo_image = T.Resize((224, 224))(pseudo_image) * 0.5 + 0.5
+            pseudo_image = self.clip_normalize(pseudo_image)
+
+            # Extract image embedding
+            clip_feat_image = self.clip_model.get_image_features(pseudo_image)
+            clip_feat_image = clip_feat_image / clip_feat_image.norm(dim=-1, keepdim=True)
+
+            # Extract text embedding
+            clip_feat_text_pos, clip_feat_text_neg = torch.split(
+                text_features, [self.num_pos_samples, self.num_neg_samples]
+            )
+
+            if contrastive:
+                clip_loss = self.contrast_criterion(
+                    clip_feat_image, torch.cat((clip_feat_text_pos, clip_feat_text_neg), dim=0),
+                    torch.cat((torch.zeros(self.num_pos_samples), torch.ones(self.num_neg_samples))).unsqueeze(0).cuda()
+                )
+            else:
+                clip_loss = 1. - (clip_feat_image @ clip_feat_text_pos.T).mean()
+
+            loss = self.clip_alpha * clip_loss
+
+        return -torch.autograd.grad(loss, x, allow_unused=True)[0]
+
+    def sample(self, original_label, class_index, batch_size=1, device=None):
+        z = torch.randn(batch_size, 4, self.latent_size, self.latent_size, device=device)
+        y = torch.tensor([original_label] * batch_size, device=device)
+
+        # classifier-free guidance
+        z = torch.cat([z, z], 0)
+        y_null = torch.tensor([1000] * batch_size, device=device)
+        y = torch.cat([y, y_null], 0)
+
+        pos_descriptions = self.descriptions[class_index]
+        pos_descriptions = [self.class_names[class_index]+' with '+description for description in pos_descriptions]
+        neg_descriptions = []
+        if self.neg_policy == 'random':
+            neg_classes = random.choices(self.class_indices, k=self.num_neg_samples)
+        elif self.neg_policy == 'similar':
+            max_indices = np.argsort(self.neighbors[class_index])[-self.num_neg_samples:]
+            neg_classes = [self.class_indices[max_index] for max_index in max_indices]
+        else:
+            neg_classes = random.choices(self.class_indices, self.neighbors[class_index], k=self.num_neg_samples)
+        for rand_index in neg_classes:
+            neg_descriptions.append(self.class_names[rand_index] + ' with ' + self.descriptions[rand_index][np.random.randint(0, 4)])
+        all_descriptions = pos_descriptions + neg_descriptions
+        text_features = self.extract_text_feature(all_descriptions)
+
+        model_kwargs = dict(
+            y=y, cfg_scale=self.cfg_scale,
+            text_features=text_features, contrastive=True
+        )
+
+        def get_samples(z):
+            samples = self.diffusion.ddim_sample_loop(
+                self.dit.forward_with_cfg, z.shape, z, clip_denoised=False,
+                model_kwargs=model_kwargs, progress=False, device=device,
+                cond_fn=self.cond_fn
+            )
+            samples, _ = samples.chunk(2, dim=0)
+            samples = self.vae.decode(samples / 0.18215).sample
+            samples = T.Resize((224, 224))(samples)
+
+            return samples
+
+        samples = get_samples(z)
+
+        return samples

var/D3HR/validation/utils/syn_utils_img2img.py

@@ -0,0 +1,134 @@
+import json
+import torch
+import random
+import numpy as np
+from torchvision import transforms as T
+from transformers import CLIPModel, AutoTokenizer
+
+from misc import prompts
+from models.scheduling_ddim import DDIMScheduler
+
+
+class SupConLoss(torch.nn.Module):
+    def __init__(self, temperature=0.05, base_temperatue=0.05):
+        super(SupConLoss, self).__init__()
+        self.temperature = temperature
+        self.base_temperature = base_temperatue
+    
+    def forward(self, image_features, text_features, text_labels):
+        logits = (image_features @ text_features.T) / self.temperature
+        logits_max, _ = torch.max(logits, dim=1, keepdim=True)
+        logits = logits - logits_max.detach()
+
+        exp_logits = torch.exp(logits) * text_labels
+        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
+        mean_log_prob_pos = ((1 - text_labels) * log_prob).sum(1) / (1 - text_labels).sum(1)
+        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+        loss = loss.mean()
+
+        return loss
+
+
+class ImageSynthesizer(object):
+    def __init__(self, args):
+        self.init_clip()
+        self.description_file = args.description_path
+        self.load_class_description()
+        self.contrast_criterion = SupConLoss()
+
+        self.prompts = prompts.prompt_templates
+        self.diffusion_steps = args.diffusion_steps
+        self.clip_alpha = args.clip_alpha
+        self.num_neg_samples = args.num_neg_samples
+        self.neg_policy = args.neg_policy
+
+    def load_class_description(self):
+        with open(self.description_file, 'r') as fp:
+            descriptions = json.load(fp)
+        self.class_names = {}
+        self.descriptions = {}
+
+        for class_index, (class_name, description) in descriptions.items():
+            self.class_names[class_index] = class_name
+            self.descriptions[class_index] = description
+        
+        self.class_indices = list(self.class_names.keys())
+        self.class_name_list = list(self.class_names.values())
+        self.neighbors = {}
+        with torch.no_grad():
+            class_name_feat = self.extract_clip_text_embed(self.class_name_list)
+            name_sims = (class_name_feat @ class_name_feat.T).cpu()
+            name_sims -= torch.eye(len(name_sims))
+            name_sims = name_sims.numpy()
+        for class_index, sim_indices in zip(self.class_indices, name_sims):
+            self.neighbors[class_index] = list(sim_indices)
+
+    def init_clip(self):
+        self.clip_model = CLIPModel.from_pretrained('laion/CLIP-ViT-L-14-laion2B-s32B-b82K').to('cuda')
+        self.clip_tokenizer = AutoTokenizer.from_pretrained('laion/CLIP-ViT-L-14-laion2B-s32B-b82K')
+        self.clip_normalize = T.Normalize(
+            mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]
+        )
+
+    def extract_clip_image_embed(self, image):
+        image = self.clip_transform(image).unsqueeze(0)
+        clip_feat = self.clip_model.encode_image(image)
+        return clip_feat
+
+    def extract_clip_text_embed(self, descriptions):
+        input_text = self.clip_tokenizer(descriptions, padding=True, return_tensors='pt').to('cuda')
+        text_features = self.clip_model.get_text_features(**input_text)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+        return text_features
+
+    def cond_fn(self, sample, **kwargs):
+        with torch.enable_grad():
+            sample = torch.nn.Parameter(sample).requires_grad_()
+            pseudo_image = self.pipe.vae.decode(sample / 0.18215, return_dict=False)[0]
+            pseudo_image = T.Resize((224, 224))(pseudo_image) * 0.5 + 0.5
+            pseudo_image = self.clip_normalize(pseudo_image)
+
+            # Extract image embedding
+            clip_feat_image = self.clip_model.get_image_features(pseudo_image)
+            clip_feat_image = clip_feat_image / clip_feat_image.norm(dim=-1, keepdim=True)
+
+            clip_loss = self.contrast_criterion(
+                clip_feat_image, self.current_desc_embeddings,
+                torch.cat((torch.zeros(5), torch.ones(len(self.current_desc_embeddings) - 5))).unsqueeze(0).cuda()
+            )
+
+            loss = self.clip_alpha * clip_loss
+
+        return -torch.autograd.grad(loss, sample, allow_unused=True)[0]
+
+    def init_img2img(self):
+        from models.pipeline_stable_unclip_img2img import StableUnCLIPImg2ImgPipeline
+        self.pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
+            'radames/stable-diffusion-2-1-unclip-img2img'
+        )
+        self.pipe.scheduler = DDIMScheduler.from_pretrained('radames/stable-diffusion-2-1-unclip-img2img', subfolder='scheduler')
+        self.pipe = self.pipe.to('cuda')
+
+    def sample_img2img(self, image, class_index):
+        class_name = self.class_names[class_index]
+        class_name = class_name.split(',')[0]
+        pos_descriptions = self.descriptions[class_index]
+        prompt = random.choice(self.prompts).format(class_name, '')
+
+        pos_descriptions = [self.class_names[class_index]+' with '+description for description in pos_descriptions]
+        neg_descriptions = []
+        if self.neg_policy == 'random':
+            neg_classes = random.choices(self.class_indices, k=self.num_neg_samples)
+        elif self.neg_policy == 'similar':
+            max_indices = np.argsort(self.neighbors[class_index])[-self.num_neg_samples:]
+            neg_classes = [self.class_indices[max_index] for max_index in max_indices]
+        else:
+            neg_classes = random.choices(self.class_indices, self.neighbors[class_index], k=self.num_neg_samples)
+        for rand_index in neg_classes:
+            neg_descriptions.append(self.class_names[rand_index] + 'with ' + self.descriptions[rand_index][np.random.randint(0, 4)])
+        self.current_desc_embeddings = self.extract_clip_text_embed(pos_descriptions + neg_descriptions)
+        new_image = self.pipe(image=image, prompt=prompt, cond_fn=self.cond_fn, num_inference_steps=self.diffusion_steps).images[0]
+
+        new_image = new_image.resize((224, 224))
+
+        return new_image