diff --git a/brushnet_nodes.py b/brushnet_nodes.py index 655202a..15a8a54 100644 --- a/brushnet_nodes.py +++ b/brushnet_nodes.py @@ -19,7 +19,7 @@ class BrushNetLoader: def INPUT_TYPES(s): return {"required": { - "brushnet": ([x[0][17:] for x in os.walk("./models/inpaint/")], ), + "brushnet": (inpaint_folders(), ), }, } @@ -30,12 +30,27 @@ class BrushNetLoader: FUNCTION = "brushnet_loading" def brushnet_loading(self, brushnet): - brushnet_path = "./models/inpaint/"+brushnet + brushnet_path = os.path.join(folder_paths.models_dir, "inpaint", brushnet) print("Brushnet... ", end="") brush_net = BrushNetModel.from_pretrained(brushnet_path, torch_dtype=torch.float16) print("loaded") return (brush_net,) + +def inpaint_folders(): + inpaint_path = os.path.join(folder_paths.models_dir, 'inpaint') + abs_list = [x[0] for x in os.walk(inpaint_path)] + + folders = [] + for x in abs_list[1:]: + remain = x + y = '' + while remain != inpaint_path: + remain, folder = os.path.split(remain) + y = os.path.join(folder, y) + folders.append(y) + return folders + class BrushNetPipeline: @@ -61,7 +76,8 @@ class BrushNetPipeline: model_path, brushnet=brushnet, torch_dtype=torch.float16, - low_cpu_mem_usage=False + low_cpu_mem_usage=False, + original_config_file=os.path.join(os.path.dirname(__file__), "v1-inference.yaml"), ) pipe.to("cuda") pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) diff --git a/v1-inference.yaml b/v1-inference.yaml new file mode 100644 index 0000000..d4effe5 --- /dev/null +++ b/v1-inference.yaml @@ -0,0 +1,70 @@ +model: + base_learning_rate: 1.0e-04 + target: ldm.models.diffusion.ddpm.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: crossattn + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: False + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 10000 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: True + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: ldm.modules.encoders.modules.FrozenCLIPEmbedder