diff --git a/example1_standard.py b/example1_standard.py index b8975bc..e0e0d04 100644 --- a/example1_standard.py +++ b/example1_standard.py @@ -32,11 +32,9 @@ from stable_diffusion_holder import StableDiffusionHolder torch.set_grad_enabled(False) #%% First let us spawn a stable diffusion holder -device = "cuda" fp_ckpt = "../stable_diffusion_models/ckpt/v2-1_768-ema-pruned.ckpt" -fp_config = 'configs/v2-inference-v.yaml' -sdh = StableDiffusionHolder(fp_ckpt, fp_config, device) +sdh = StableDiffusionHolder(fp_ckpt) #%% Next let's set up all parameters diff --git a/example2_inpaint.py b/example2_inpaint.py index e5f9d0e..e2de40c 100644 --- a/example2_inpaint.py +++ b/example2_inpaint.py @@ -34,11 +34,8 @@ from stable_diffusion_holder import StableDiffusionHolder torch.set_grad_enabled(False) #%% First let us spawn a stable diffusion holder -device = "cuda" fp_ckpt= "../stable_diffusion_models/ckpt/512-inpainting-ema.ckpt" -fp_config = '../stablediffusion/configs//stable-diffusion/v2-inpainting-inference.yaml' -sdh = StableDiffusionHolder(fp_ckpt, fp_config, device) - +sdh = StableDiffusionHolder(fp_ckpt) #%% Let's first make a source image and mask. quality = 'medium' diff --git a/example3_multitrans.py b/example3_multitrans.py index 2c09a85..451bdec 100644 --- a/example3_multitrans.py +++ b/example3_multitrans.py @@ -31,10 +31,8 @@ from stable_diffusion_holder import StableDiffusionHolder torch.set_grad_enabled(False) #%% First let us spawn a stable diffusion holder -device = "cuda" fp_ckpt = "../stable_diffusion_models/ckpt/v2-1_768-ema-pruned.ckpt" -fp_config = 'configs/v2-inference-v.yaml' -sdh = StableDiffusionHolder(fp_ckpt, fp_config, device) +sdh = StableDiffusionHolder(fp_ckpt) #%% Let's setup the multi transition diff --git a/example4_upscaling.py b/example4_upscaling.py index 74888ba..5222940 100644 --- a/example4_upscaling.py +++ b/example4_upscaling.py @@ -45,17 +45,15 @@ depth_strength_lores = 0.5 device = "cuda" fp_ckpt_lores = "../stable_diffusion_models/ckpt/v2-1_512-ema-pruned.ckpt" -fp_config_lores = 'configs/v2-inference.yaml' #%% Define vars for high-resoltion pass fp_ckpt_hires = "../stable_diffusion_models/ckpt/x4-upscaler-ema.ckpt" -fp_config_hires = 'configs/x4-upscaling.yaml' depth_strength_hires = 0.65 num_inference_steps_hires = 100 nmb_branches_final_hires = 6 #%% Run low-res pass -sdh = StableDiffusionHolder(fp_ckpt_lores, fp_config_lores, device) +sdh = StableDiffusionHolder(fp_ckpt_lores) lb = LatentBlending(sdh) lb.set_prompt1(prompt1) lb.set_prompt2(prompt2) @@ -64,6 +62,6 @@ lb.set_height(height) lb.run_upscaling_step1(dp_img, depth_strength_lores, num_inference_steps_lores, nmb_branches_final_lores, fixed_seeds) #%% Run high-res pass -sdh = StableDiffusionHolder(fp_ckpt_hires, fp_config_hires) +sdh = StableDiffusionHolder(fp_ckpt_hires) lb = LatentBlending(sdh) lb.run_upscaling_step2(dp_img, depth_strength_hires, num_inference_steps_hires, nmb_branches_final_hires) \ No newline at end of file diff --git a/stable_diffusion_holder.py b/stable_diffusion_holder.py index d9b1972..05cbcaf 100644 --- a/stable_diffusion_holder.py +++ b/stable_diffusion_holder.py @@ -163,8 +163,27 @@ class StableDiffusionHolder: """ assert os.path.isfile(fp_ckpt), f"Your model checkpoint file does not exist: {fp_ckpt}" - assert os.path.isfile(fp_config), f"Your config file does not exist: {fp_config}" self.fp_ckpt = fp_ckpt + + # Auto init the config? + if fp_config is None: + fn_ckpt = os.path.basename(fp_ckpt) + if 'depth' in fn_ckpt: + fp_config = 'configs/v2-midas-inference.yaml' + elif 'inpain' in fn_ckpt: + fp_config = 'configs/v2-inpainting-inference.yaml' + elif 'upscaler' in fn_ckpt: + fp_config = 'configs/x4-upscaling.yaml' + elif '512' in fn_ckpt: + fp_config = 'configs/v2-inference.yaml' + elif '768'in fn_ckpt: + fp_config = 'configs/v2-inference-v.yaml' + else: + raise ValueError("auto detect of config failed. please specify fp_config manually!") + + + assert os.path.isfile(fp_config), f"Your config file does not exist: {fp_config}" + config = OmegaConf.load(fp_config) @@ -174,6 +193,17 @@ class StableDiffusionHolder: self.model = self.model.to(self.device) self.sampler = DDIMSampler(self.model) + + def init_auto_res(self): + r"""Automatically set the resolution to the one used in training. + """ + if '768' in self.fp_ckpt: + self.height = 768 + self.width = 768 + else: + self.height = 512 + self.width = 512 + def set_negative_prompt(self, negative_prompt): r"""Set the negative prompt. Currenty only one negative prompt is supported """ @@ -185,17 +215,6 @@ class StableDiffusionHolder: if len(self.negative_prompt) > 1: self.negative_prompt = [self.negative_prompt[0]] - - def init_auto_res(self): - r"""Automatically set the resolution to the one used in training. - """ - if '768' in self.fp_ckpt: - self.height = 768 - self.width = 768 - else: - self.height = 512 - self.width = 512 - def init_inpainting( self, @@ -571,10 +590,12 @@ if __name__ == "__main__": # fp_config = '../stablediffusion/configs//stable-diffusion/v2-inpainting-inference.yaml' fp_ckpt = "../stable_diffusion_models/ckpt/v2-1_768-ema-pruned.ckpt" - fp_config = 'configs/v2-inference-v.yaml' + # fp_config = 'configs/v2-inference-v.yaml' - self = StableDiffusionHolder(fp_ckpt, fp_config, num_inference_steps) + self = StableDiffusionHolder(fp_ckpt, num_inference_steps=num_inference_steps) + + xxx #%% self.width = 1536