(original_unet_config)
| 776 | |
| 777 | |
| 778 | def superres_create_unet_diffusers_config(original_unet_config): |
| 779 | attention_resolutions = parse_list(original_unet_config["attention_resolutions"]) |
| 780 | attention_resolutions = [original_unet_config["image_size"] // int(res) for res in attention_resolutions] |
| 781 | |
| 782 | channel_mult = parse_list(original_unet_config["channel_mult"]) |
| 783 | block_out_channels = [original_unet_config["model_channels"] * mult for mult in channel_mult] |
| 784 | |
| 785 | down_block_types = [] |
| 786 | resolution = 1 |
| 787 | |
| 788 | for i in range(len(block_out_channels)): |
| 789 | if resolution in attention_resolutions: |
| 790 | block_type = "SimpleCrossAttnDownBlock2D" |
| 791 | elif original_unet_config["resblock_updown"]: |
| 792 | block_type = "ResnetDownsampleBlock2D" |
| 793 | else: |
| 794 | block_type = "DownBlock2D" |
| 795 | |
| 796 | down_block_types.append(block_type) |
| 797 | |
| 798 | if i != len(block_out_channels) - 1: |
| 799 | resolution *= 2 |
| 800 | |
| 801 | up_block_types = [] |
| 802 | for i in range(len(block_out_channels)): |
| 803 | if resolution in attention_resolutions: |
| 804 | block_type = "SimpleCrossAttnUpBlock2D" |
| 805 | elif original_unet_config["resblock_updown"]: |
| 806 | block_type = "ResnetUpsampleBlock2D" |
| 807 | else: |
| 808 | block_type = "UpBlock2D" |
| 809 | up_block_types.append(block_type) |
| 810 | resolution //= 2 |
| 811 | |
| 812 | head_dim = original_unet_config["num_head_channels"] |
| 813 | use_linear_projection = ( |
| 814 | original_unet_config["use_linear_in_transformer"] |
| 815 | if "use_linear_in_transformer" in original_unet_config |
| 816 | else False |
| 817 | ) |
| 818 | if use_linear_projection: |
| 819 | # stable diffusion 2-base-512 and 2-768 |
| 820 | if head_dim is None: |
| 821 | head_dim = [5, 10, 20, 20] |
| 822 | |
| 823 | class_embed_type = None |
| 824 | projection_class_embeddings_input_dim = None |
| 825 | |
| 826 | if "num_classes" in original_unet_config: |
| 827 | if original_unet_config["num_classes"] == "sequential": |
| 828 | class_embed_type = "projection" |
| 829 | assert "adm_in_channels" in original_unet_config |
| 830 | projection_class_embeddings_input_dim = original_unet_config["adm_in_channels"] |
| 831 | else: |
| 832 | raise NotImplementedError( |
| 833 | f"Unknown conditional unet num_classes config: {original_unet_config['num_classes']}" |
| 834 | ) |
| 835 |
no test coverage detected
searching dependent graphs…