ai-art

 1"""
 2Checks to make sure all our files and folders exist in /data. We need in the
 3data folder:
 4
 5- /data/models/
 6- /data/models/vqgan_imagenet_f16_16384.json
 7- /data/models/vqgan_imagenet_f16_16384.ckpt
 8- /data/outputs/
 9- /data/outputs/steps/
10- /data/config.json
11"""
12import requests
13import os
14
15vqgan_imagenet_f16_16384_ckpt_url = "https://heibox.uni-heidelberg.de/f/867b05fc8c4841768640/?dl=1"
16
17vqgan_imagenet_f16_16384_json = """{
18  "params": {
19    "embed_dim": 256,
20    "n_embed": 16384,
21    "ddconfig": {
22      "double_z": false,
23      "z_channels": 256,
24      "resolution": 256,
25      "in_channels": 3,
26      "out_ch": 3,
27      "ch": 128,
28      "ch_mult": [1, 1, 2, 2, 4],
29      "num_res_blocks": 2,
30      "attn_resolutions": [16],
31      "dropout": 0.0
32    },
33    "lossconfig": {
34      "params": {
35        "disc_conditional": false,
36        "disc_in_channels": 3,
37        "disc_start": 0,
38        "disc_weight": 0.75,
39        "disc_num_layers": 2,
40        "codebook_weight": 1.0
41      }
42    }
43  }
44}
45"""
46
47config_json = """{
48    "prompts": ["space", "fractal"],
49    "init_image": "",
50    "size": [256, 256],
51    "max_iterations": 250,
52    "save_freq": 50
53}
54"""
55
56
57def check_files_and_folders():
58    print("Checking that you have all the files and folders required...")
59
60    # check that models folder exists, if not create it
61    if not os.path.exists("/data/models"):
62        print("Creating models folder...")
63        os.makedirs("/data/models")
64
65    # check that outputs folder exists, if not create it
66    if not os.path.exists("/data/outputs"):
67        os.makedirs("/data/outputs")
68        os.makedirs("/data/outputs/steps")
69
70    # check that config.json exists, if not create it
71    if not os.path.exists("/data/config.json"):
72        print("Creating config.json...")
73        with open("/data/config.json", "w") as f:
74            f.write(config_json)
75
76    # check that vqgan_imagenet_f16_16384.json exists, if not create it
77    if not os.path.exists("/data/models/vqgan_imagenet_f16_16384.json"):
78        print("Creating vqgan_imagenet_f16_16384.json...")
79        with open("/data/models/vqgan_imagenet_f16_16384.json", "w") as f:
80            f.write(vqgan_imagenet_f16_16384_json)
81            f.close()
82
83    # check that vqgan_imagenet_f16_16384.ckpt exists, if not download and
84    # write in chunks since it's a large file
85    if not os.path.exists("/data/models/vqgan_imagenet_f16_16384.ckpt"):
86        print("Downloading vqgan_imagenet_f16_16384.ckpt...")
87        with open("/data/models/vqgan_imagenet_f16_16384.ckpt", "wb") as f:
88            r = requests.get(vqgan_imagenet_f16_16384_ckpt_url)
89            for chunk in r.iter_content(chunk_size=1024):
90                if chunk:
91                    f.write(chunk)
92                    f.flush()
93            f.close()