init

net/UNet.py

@@ -422,99 +422,3 @@ class UNet(nn.Module):
         # print("decoder:")
         # print(decoder_out.shape)
         return decoder_out
-
-
-if __name__ == '__main__':
-    import cv2, os
-
-
-    def modelSave(model, save_path, save_name):
-        if not os.path.exists(save_path):
-            os.mkdir(save_path)
-        torch.save(model.state_dict(), os.path.join(save_path, save_name))
-
-
-    def merge_images(images: np.ndarray):
-        """
-        合并图像
-        :param images: 图像数组
-        :return: 合并后的图像数组
-        """
-        n, h, w, c = images.shape
-        nn = int(np.ceil(n ** 0.5))
-        merged_image = np.zeros((h * nn, w * nn, 3), dtype=images.dtype)
-        for i in range(n):
-            row = i // nn
-            col = i % nn
-            merged_image[row * h:(row + 1) * h, col * w:(col + 1) * w, :] = images[i]
-
-        merged_image = np.clip(merged_image, 0, 255)
-        merged_image = np.array(merged_image, dtype=np.uint8)
-        return merged_image
-
-
-    # 320,448,576,832
-    config = {  # 模型结构相关
-        "en_out_c": (256, 256, 256, 320, 320, 320, 576, 576, 576, 704, 704, 704),
-        "en_down": (0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0),
-        "en_skip": (0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1),
-        "en_att_heads": (8, 8, 8, 0, 8, 8, 0, 8, 8, 0, 8, 8),
-        "de_out_c": (704, 576, 576, 576, 320, 320, 320, 256, 256, 256, 256),
-        "de_up": ("none", "subpix", "none", "none", "subpix", "none", "none", "subpix", "none", "none", "none"),
-        "de_skip": (1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0),
-        "de_att_heads": (8, 8, 0, 8, 8, 0, 8, 8, 0, 8, 8),  # skip的地方不做self-attention
-        "t_out_c": 256,
-        "vae_c": 4,
-        "block_deep": 3,
-    }
-    device = "cuda"
-    total_step = 1000
-
-    unet = UNet(config["en_out_c"], config["en_down"], config["en_skip"], config["en_att_heads"],
-                config["de_out_c"], config["de_up"], config["de_skip"], config["de_att_heads"],
-                config["t_out_c"], config["vae_c"], config["block_deep"]).to(device)
-
-    print("总参数", sum(i.numel() for i in unet.parameters()) / 10000, "单位:万")
-    print("encoder", sum(i.numel() for i in unet.encoder.parameters()) / 10000, "单位:万")
-    print("decoder", sum(i.numel() for i in unet.decoder.parameters()) / 10000, "单位:万")
-    print("t", sum(i.numel() for i in unet.t_encoder.parameters()) / 10000, "单位:万")
-
-    batch_size = 2
-    x = np.random.random((batch_size, config["vae_c"], 32, 32))
-    t = np.random.uniform(1, total_step + 0.9999, size=(batch_size, 1))
-    t = np.array(t, dtype=np.int16)
-    t = t / total_step
-
-    with torch.no_grad():
-        x = torch.Tensor(x).to(device)
-        t = torch.Tensor(t).to(device)
-        y = unet(x, t)
-        print(y.shape)
-
-        z = y[0].cpu().numpy()
-        # z = (z - np.mean(z)) / (np.max(z) - np.min(z))
-        z = np.clip(np.asarray((z + 1) * 127.5), 0, 255)
-        z = np.asarray(z, dtype=np.uint8)
-
-        z = [np.tile(z[ii, :, :, np.newaxis], (1, 1, 3)) for ii in range(z.shape[0])]
-        noise = merge_images(np.array(z))
-
-        noise = cv2.resize(noise, None, fx=2, fy=2)
-        cv2.imshow("noise", noise)
-        cv2.waitKey(0)
-
-    # modelSave(unet, "./", "test.pth")
-    # 导出为onnx格式
-    torch.onnx.export(
-        unet,
-        (x, t),
-        'unet.onnx',
-        export_params=True,
-        opset_version=12,
-    )
-    import onnx
-
-    # 增加维度信息
-    model_file = 'unet.onnx'
-    onnx_model = onnx.load(model_file)
-    onnx.save(onnx.shape_inference.infer_shapes(onnx_model), model_file)

requirements.txt

@@ -1,186 +1,9 @@
-absl-py==1.3.0
-addict==2.4.0
-aiofiles==23.1.0
-aiohttp==3.8.3
-aiosignal==1.3.1
-aliyun-python-sdk-core==2.13.36
-aliyun-python-sdk-kms==2.16.0
-altair==4.2.0
-anyio==3.6.2
-appdirs==1.4.4
-asttokens==2.3.0
-async-timeout==4.0.2
-attrs==22.1.0
-audioread==3.0.0
-backcall==0.2.0
-certifi==2022.12.7
-cffi==1.15.1
-charset-normalizer==2.1.1
-chumpy==0.70
-click==8.1.3
-clip==1.0
-colorama==0.4.6
-commonmark==0.9.1
-contourpy==1.0.6
-cpm-kernels==1.0.11
-crcmod==1.7
-cryptography==39.0.2
-cycler==0.11.0
-Cython==0.29.32
-datasets==2.8.0
-decorator==5.1.1
-decord==0.6.0
-diffusers==0.20.1
-dill==0.3.6
-docker-pycreds==0.4.0
-einops==0.6.0
-entrypoints==0.4
-exceptiongroup==1.1.3
-executing==1.2.0
-fastapi==0.88.0
-ffmpy==0.3.0
-filelock==3.8.2
-Flask==2.0.2
-Flask-Cors==3.0.10
-fonttools==4.38.0
-frozenlist==1.3.3
-fsspec==2022.11.0
-ftfy==6.1.1
-gast==0.5.3
-gitdb==4.0.10
-GitPython==3.1.32
-gradio==3.39.0
-gradio_client==0.3.0
-h11==0.14.0
-httpcore==0.16.2
-httpx==0.23.1
+gradio
+gradio_client
 huggingface-hub==0.16.4
-icetk==0.0.4
-idna==3.4
-importlib-metadata==5.2.0
-ipython==8.15.0
-itsdangerous==2.1.2
-jedi==0.19.0
-Jinja2==3.1.2
-jmespath==0.10.0
-joblib==1.2.0
-json-tricks==3.16.1
-jsonplus==0.8.0
-jsonschema==4.17.3
-kiwisolver==1.4.4
-lazy_loader==0.1
-librosa==0.10.0
-linkify-it-py==1.0.3
-lion-pytorch==0.1.2
-llvmlite==0.39.1
-loguru==0.6.0
-Markdown==3.4.1
-markdown-it-py==2.1.0
-MarkupSafe==2.1.1
-matplotlib==3.6.2
-matplotlib-inline==0.1.6
-mdit-py-plugins==0.3.3
-mdurl==0.1.2
-mediapipe==0.8.11
-mmcv-full==1.7.0
-mmdet==2.26.0
-model-index==0.1.11
-modelscope==1.3.2
-mpmath==1.2.1
-msgpack==1.0.4
-multidict==6.0.3
-multiprocess==0.70.14
-munkres==1.1.4
-networkx==3.0
-numba==0.56.4
 numpy==1.23.4
-onnx==1.14.1
-opencv-contrib-python==4.5.5.64
-opencv-python==4.5.5.64
-openmim==0.3.3
-ordered-set==4.1.0
-orjson==3.8.3
-oss2==2.16.0
-packaging==21.3
-pandas==1.5.2
-parso==0.8.3
-pathtools==0.1.2
-pickleshare==0.7.5
-Pillow==9.2.0
-pip==23.1.2
-platformdirs==3.1.0
-plotly==5.11.0
-pooch==1.7.0
-prodigyopt==1.0
-prompt-toolkit==3.0.39
-protobuf==4.24.2
-psutil==5.9.5
-pure-eval==0.2.2
-pyarrow==11.0.0
-pycocotools==2.0.6
-pycparser==2.21
-pycryptodome==3.16.0
-pydantic==1.10.2
-pydub==0.25.1
-Pygments==2.13.0
-pyparsing==3.0.9
-pyrsistent==0.19.2
-python-dateutil==2.8.2
-python-multipart==0.0.5
-pytorch-fid==0.3.0
-pytz==2022.6
-PyYAML==6.0
-regex==2022.10.31
-requests==2.28.1
-responses==0.18.0
-rfc3986==1.5.0
-rich==12.6.0
-safetensors==0.3.3
-scikit-learn==1.2.1
-scipy==1.9.3
-semantic-version==2.10.0
-sentencepiece==0.1.97
-sentry-sdk==1.28.0
-setproctitle==1.3.2
-setuptools==65.5.0
-simplejson==3.18.3
-six==1.16.0
-smmap==5.0.0
-sniffio==1.3.0
-sortedcontainers==2.4.0
-soundfile==0.12.1
-soxr==0.3.4
-stack-data==0.6.2
-starlette==0.22.0
-sympy==1.11.1
-tabulate==0.9.0
-tenacity==8.1.0
-terminaltables==3.1.10
-threadpoolctl==3.1.0
-timm==0.4.9
-tokenizers==0.13.2
-toolz==0.12.0
 torch==2.0.0+cu117
 torchaudio==2.0.1+cu117
 torchinfo==1.7.1
 torchvision==0.15.1+cu117
 tqdm==4.64.1
-traitlets==5.9.0
-transformers==4.26.1
-typing_extensions==4.4.0
-uc-micro-py==1.0.1
-unicodedata2==15.0.0
-urllib3==1.26.12
-uvicorn==0.20.0
-wandb==0.15.5
-wcwidth==0.2.5
-websockets==10.4
-Werkzeug==2.2.2
-wheel==0.37.1
-win32-setctime==1.1.0
-wincertstore==0.2
-xtcocotools==1.12
-xxhash==3.2.0
-yapf==0.32.0
-yarl==1.8.2
-zipp==3.11.0