From 386fce4c0a85c30e580ad414d252ff129f1ccb86 Mon Sep 17 00:00:00 2001 From: Anatoly Belikov Date: Fri, 11 Aug 2023 13:35:34 +0300 Subject: [PATCH 1/3] merge --- examples/im2im.py | 2 +- examples/prompt2im_hypernet.py | 4 ++-- multigen/pipes.py | 43 ++++++++++++++++++++-------------- multigen/sessions.py | 9 +++---- 4 files changed, 33 insertions(+), 25 deletions(-) diff --git a/examples/im2im.py b/examples/im2im.py index cc5aeeb..b075c6e 100755 --- a/examples/im2im.py +++ b/examples/im2im.py @@ -10,7 +10,7 @@ prompt = ["bioinformatics lab with flasks and exotic flowers", "happy vibrant", "green colors", "artwork", "high tech"] -#pipe = CIm2ImPipe(model_dir+model_id) +#pipe = ControlNetIm2ImPipe(model_dir+model_id) pipe = Im2ImPipe(model_dir+model_id) pipe.setup("./_projects/biolab/00000.png", strength=0.95, scheduler=get_diffusion_scheduler_names()[0]) diff --git a/examples/prompt2im_hypernet.py b/examples/prompt2im_hypernet.py index c580d70..af8450b 100755 --- a/examples/prompt2im_hypernet.py +++ b/examples/prompt2im_hypernet.py @@ -10,8 +10,8 @@ path = "./models/icb_diffusers_final/" path_hypernet = './models/hypernetworks/digitalben.pt' -pipe = Prompt2ImPipe(path, lpw=True, scheduler='DPMSolverMultistepScheduler') -pipe.setup(width=512, height=512, guidance_scale=5.5, clip_skip=1) +pipe = Prompt2ImPipe(path, lpw=True) +pipe.setup(width=512, height=512, guidance_scale=5.5, clip_skip=1, scheduler='DPMSolverMultistepScheduler') pipe.add_hypernet(path_hypernet, 0.65) # pipe.clear_hypernets() diff --git a/multigen/pipes.py b/multigen/pipes.py index 5ec8e51..5d484ef 100755 --- a/multigen/pipes.py +++ b/multigen/pipes.py @@ -1,4 +1,6 @@ import importlib + +import PIL import torch from PIL import Image @@ -124,6 +126,7 @@ def setup(self, steps=50, **args): # TODO? add scheduler to config? self.try_set_scheduler(dict(scheduler=args['scheduler'])) + class Prompt2ImPipe(BasePipe): def __init__(self, model_id: str, @@ -143,8 +146,8 @@ def setup(self, width=768, height=768, guidance_scale=7.5, **args): "guidance_scale": guidance_scale }) - def gen(self, inputs): - inputs = {**inputs} + def gen(self, inputs: dict) -> PIL.Image.Image: + inputs = inputs.copy() inputs.update(self.pipe_params) # allow for scheduler overwrite self.try_set_scheduler(inputs) @@ -158,18 +161,17 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionImg2ImgPipeline] = No super().__init__(model_id=model_id, sd_pipe_class=StableDiffusionImg2ImgPipeline, pipe=pipe, **args) self._input_image = None - def setup(self, fimage, image=None, strength=0.75, gscale=7.5, scale=None, **args): + def setup(self, img_path, image=None, strength=0.75, gscale=7.5, scale=None, **args): super().setup(**args) - self.fname = fimage - self._input_image = Image.open(fimage).convert("RGB") if image is None else image + self.fname = img_path + self._input_image = Image.open(img_path).convert("RGB") if image is None else image if scale is not None: if not isinstance(scale, list): scale = [8 * (int(self._input_image.size[i] * scale) // 8) for i in range(2)] self._input_image = self._input_image.resize((scale[0], scale[1])) self.pipe_params.update({ "strength": strength, - "guidance_scale": gscale - }) + "guidance_scale": gscale}) def get_config(self): cfg = super().get_config() @@ -179,8 +181,8 @@ def get_config(self): cfg.update(self.pipe_params) return cfg - def gen(self, inputs): - inputs = {**inputs} + def gen(self, inputs: dict) -> PIL.Image.Image: + inputs = inputs.copy() inputs.update(self.pipe_params) inputs.update({"image": self._input_image}) self.try_set_scheduler(inputs) @@ -189,6 +191,9 @@ def gen(self, inputs): class Cond2ImPipe(BasePipe): + """ + Base class for ControlNet pipelines + """ # TODO: set path cpath = "./models-cn/" @@ -218,20 +223,20 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] = self.ctypes = ctypes self._condition_image = None dtype = torch.float16 if 'torch_type' not in args else args['torch_type'] - cnets = [ControlNetModel.from_pretrained(CIm2ImPipe.cpath+CIm2ImPipe.cmodels[c], torch_dtype=dtype) for c in ctypes] + cnets = [ControlNetModel.from_pretrained(Cond2ImPipe.cpath+Cond2ImPipe.cmodels[c], torch_dtype=dtype) for c in ctypes] super().__init__(sd_pipe_class=StableDiffusionControlNetPipeline, model_id=model_id, pipe=pipe, controlnet=cnets, **args) # FIXME: do we need to setup this specific scheduler here? # should we pass its name in setup to super? self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config) - def setup(self, fimage, width=None, height=None, image=None, cscales=None, guess_mode=False, **args): + def setup(self, img_path, width=None, height=None, image=None, cscales=None, guess_mode=False, **args): super().setup(**args) # TODO: allow multiple input images for multiple control nets - self.fname = fimage - image = Image.open(fimage) if image is None else image + self.fname = img_path + image = Image.open(img_path) if image is None else image self._condition_image = [image] if cscales is None: - cscales = [CIm2ImPipe.cscalem[c] for c in self.ctypes] + cscales = [Cond2ImPipe.cscalem[c] for c in self.ctypes] self.pipe_params.update({ "width": image.size[0] if width is None else width, "height": image.size[1] if height is None else height, @@ -257,8 +262,10 @@ def gen(self, inputs): return image -class CIm2ImPipe(Cond2ImPipe): - +class ControlNetIm2ImPipe(Cond2ImPipe): + """ + ControlNet pipeline with conditional image generation + """ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] = None, ctypes=["soft"], **args): super().__init__(model_id=model_id, pipe=pipe, ctypes=ctypes, **args) @@ -280,8 +287,8 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] = self.dprocessor = DPTImageProcessor.from_pretrained("./models-other/dpt-large") self.dmodel = DPTForDepthEstimation.from_pretrained("./models-other/dpt-large") - def setup(self, fimage, width=None, height=None, image=None, cscales=None, guess_mode=False, **args): - super().setup(fimage, width, height, image, cscales, guess_mode, **args) + def setup(self, img_path, width=None, height=None, image=None, cscales=None, guess_mode=False, **args): + super().setup(img_path, width, height, image, cscales, guess_mode, **args) # Additionally process the input image # REM: CIm2ImPipe expects only one image, which can be the base for multiple control images self._condition_image = self._proc_cimg(np.asarray(self._condition_image[0])) diff --git a/multigen/sessions.py b/multigen/sessions.py index 540667f..67dcb71 100755 --- a/multigen/sessions.py +++ b/multigen/sessions.py @@ -3,15 +3,16 @@ import json from . import util from .prompting import Cfgen +from .pipes import BasePipe class GenSession: - def __init__(self, session_dir, pipe, config: Cfgen, name_prefix=""): + def __init__(self, session_dir: str, pipe: BasePipe, config: Cfgen, name_prefix=""): self.session_dir = session_dir - self.pipe = pipe - self.model_id = pipe._model_id - self.confg = config + self.pipe: BasePipe = pipe + self.model_id: str = pipe._model_id + self.confg: Cfgen = config self.last_conf = None self.name_prefix = name_prefix From 0d378101c84d5e8ce46872d6edd529ccf4d0535f Mon Sep 17 00:00:00 2001 From: Anatoly Belikov Date: Fri, 11 Aug 2023 18:36:43 +0300 Subject: [PATCH 2/3] better naming --- multigen/pipes.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/multigen/pipes.py b/multigen/pipes.py index 5d484ef..801ac36 100755 --- a/multigen/pipes.py +++ b/multigen/pipes.py @@ -197,6 +197,7 @@ class Cond2ImPipe(BasePipe): # TODO: set path cpath = "./models-cn/" + cmodels = { "canny": "sd-controlnet-canny", "pose": "control_v11p_sd15_openpose", @@ -206,7 +207,8 @@ class Cond2ImPipe(BasePipe): "depth": "control_v11f1p_sd15_depth", "inpaint": "control_v11p_sd15_inpaint" } - cscalem = { + + default_cscales = { "canny": 0.75, "pose": 1.0, "ip2p": 0.5, @@ -220,7 +222,7 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] = ctypes=["soft"], **args): if not isinstance(ctypes, list): ctypes = [ctypes] - self.ctypes = ctypes + self.control_types = ctypes self._condition_image = None dtype = torch.float16 if 'torch_type' not in args else args['torch_type'] cnets = [ControlNetModel.from_pretrained(Cond2ImPipe.cpath+Cond2ImPipe.cmodels[c], torch_dtype=dtype) for c in ctypes] @@ -236,7 +238,7 @@ def setup(self, img_path, width=None, height=None, image=None, cscales=None, gue image = Image.open(img_path) if image is None else image self._condition_image = [image] if cscales is None: - cscales = [Cond2ImPipe.cscalem[c] for c in self.ctypes] + cscales = [self.default_cscales[c] for c in self.control_types] self.pipe_params.update({ "width": image.size[0] if width is None else width, "height": image.size[1] if height is None else height, @@ -249,7 +251,7 @@ def get_config(self): cfg = super().get_config() cfg.update({ "source_image": self.fname, - "control_type": self.ctypes + "control_type": self.control_types }) cfg.update(self.pipe_params) return cfg @@ -262,7 +264,7 @@ def gen(self, inputs): return image -class ControlNetIm2ImPipe(Cond2ImPipe): +class ControlNet2ImPipe(Cond2ImPipe): """ ControlNet pipeline with conditional image generation """ @@ -295,7 +297,7 @@ def setup(self, img_path, width=None, height=None, image=None, cscales=None, gue def _proc_cimg(self, oriImg): condition_image = [] - for c in self.ctypes: + for c in self.control_types: if c == "canny": image = canny_processor(oriImg) condition_image += [Image.fromarray(image)] From c41e44f5e20ac99bdcbaa92b1d66b51e05e85045 Mon Sep 17 00:00:00 2001 From: Anatoly Belikov Date: Sun, 13 Aug 2023 13:35:38 +0300 Subject: [PATCH 3/3] remove Pipe from names of BasePipe child classes --- examples/cond2im.py | 4 ++-- examples/im2im.py | 6 +++--- examples/inp2im.py | 4 ++-- examples/prompt2im.py | 4 ++-- examples/prompt2im_hypernet.py | 4 ++-- multigen/__init__.py | 2 +- multigen/pipes.py | 16 ++++++++-------- tests/pipe_test.py | 6 +++--- 8 files changed, 23 insertions(+), 23 deletions(-) diff --git a/examples/cond2im.py b/examples/cond2im.py index e14a436..28c88bd 100644 --- a/examples/cond2im.py +++ b/examples/cond2im.py @@ -1,6 +1,6 @@ from multigen.prompting import Cfgen from multigen.sessions import GenSession -from multigen.pipes import Cond2ImPipe +from multigen.pipes import Cond2Im nprompt = "monochrome, lowres, bad anatomy, worst quality, low quality" @@ -9,7 +9,7 @@ model_dir = "./models-sd/" model_id = "icbinp" -pipe = Cond2ImPipe(model_dir+model_id, ctypes=["pose"]) +pipe = Cond2Im(model_dir + model_id, ctypes=["pose"]) pipe.setup("./pose6.jpeg", width=768, height=768) gs = GenSession("./_projects/cnet", pipe, Cfgen(prompt, nprompt)) gs.gen_sess(add_count=5) diff --git a/examples/im2im.py b/examples/im2im.py index b075c6e..e33a7ef 100755 --- a/examples/im2im.py +++ b/examples/im2im.py @@ -1,6 +1,6 @@ from multigen.prompting import Cfgen from multigen.sessions import GenSession -from multigen.pipes import Im2ImPipe, CIm2ImPipe, get_diffusion_scheduler_names +from multigen.pipes import Im2Im, ControlNet2Im, get_diffusion_scheduler_names model_dir = "./models-sd/" @@ -10,9 +10,9 @@ prompt = ["bioinformatics lab with flasks and exotic flowers", "happy vibrant", "green colors", "artwork", "high tech"] -#pipe = ControlNetIm2ImPipe(model_dir+model_id) +#pipe = ControlNet2Im(model_dir+model_id) -pipe = Im2ImPipe(model_dir+model_id) +pipe = Im2Im(model_dir+model_id) pipe.setup("./_projects/biolab/00000.png", strength=0.95, scheduler=get_diffusion_scheduler_names()[0]) gs = GenSession("./_projects/biolab/modified/", pipe, Cfgen(prompt, nprompt)) gs.gen_sess(add_count=10) diff --git a/examples/inp2im.py b/examples/inp2im.py index 2805020..72f6d92 100644 --- a/examples/inp2im.py +++ b/examples/inp2im.py @@ -1,6 +1,6 @@ from multigen.prompting import Cfgen from multigen.sessions import GenSession -from multigen.pipes import InpaintingPipe +from multigen.pipes import Inpainting from PIL import Image @@ -12,7 +12,7 @@ "happy vibrant", "green colors", "artwork", "high tech"] mask_image = Image.open("mask_up.png").resize((768, 768)) -pipe = InpaintingPipe(model_dir+model_id) +pipe = Inpainting(model_dir+model_id) pipe.setup("./_projects/biolab/00000.png", mask_image) gs = GenSession("./_projects/biolab/inpaint/", pipe, Cfgen(prompt, nprompt)) gs.gen_sess(add_count=5) diff --git a/examples/prompt2im.py b/examples/prompt2im.py index 5517187..63ef7e2 100755 --- a/examples/prompt2im.py +++ b/examples/prompt2im.py @@ -1,6 +1,6 @@ from multigen.prompting import Cfgen from multigen.sessions import GenSession -from multigen.pipes import Prompt2ImPipe +from multigen.pipes import Prompt2Im model_dir = "./models-sd/" model_id = "icbinp" @@ -15,7 +15,7 @@ ["surrealism", "impressionism", "high tech", "cyberpunk"]] -pipe = Prompt2ImPipe(model_dir+model_id) +pipe = Prompt2Im(model_dir + model_id) pipe.setup(width=768, height=768) gs = GenSession("./_projects/biolab", pipe, Cfgen(prompt, nprompt)) gs.gen_sess(add_count=10) diff --git a/examples/prompt2im_hypernet.py b/examples/prompt2im_hypernet.py index af8450b..8f5fdde 100755 --- a/examples/prompt2im_hypernet.py +++ b/examples/prompt2im_hypernet.py @@ -1,6 +1,6 @@ from multigen.prompting import Cfgen from multigen.sessions import GenSession -from multigen.pipes import Prompt2ImPipe +from multigen.pipes import Prompt2Im prompt = "photo of a (digitalben:1.1) farmer, man is on a farm next to his horse, 24mm, 4k textures, soft cinematic light, RAW photo, photorealism, photorealistic, highly detailed, sharp focus, soothing tones, insane details, intricate details, hyperdetailed, low contrast, soft cinematic light, dim colors, exposure blend, hdr, faded" @@ -10,7 +10,7 @@ path = "./models/icb_diffusers_final/" path_hypernet = './models/hypernetworks/digitalben.pt' -pipe = Prompt2ImPipe(path, lpw=True) +pipe = Prompt2Im(path, lpw=True) pipe.setup(width=512, height=512, guidance_scale=5.5, clip_skip=1, scheduler='DPMSolverMultistepScheduler') pipe.add_hypernet(path_hypernet, 0.65) # pipe.clear_hypernets() diff --git a/multigen/__init__.py b/multigen/__init__.py index cae42d9..e1d0d8e 100755 --- a/multigen/__init__.py +++ b/multigen/__init__.py @@ -1,3 +1,3 @@ -from .pipes import Prompt2ImPipe, Im2ImPipe +from .pipes import Prompt2Im, Im2Im, ControlNet2Im from .sessions import GenSession from .prompting import Cfgen \ No newline at end of file diff --git a/multigen/pipes.py b/multigen/pipes.py index 801ac36..c3648d4 100755 --- a/multigen/pipes.py +++ b/multigen/pipes.py @@ -127,7 +127,7 @@ def setup(self, steps=50, **args): self.try_set_scheduler(dict(scheduler=args['scheduler'])) -class Prompt2ImPipe(BasePipe): +class Prompt2Im(BasePipe): def __init__(self, model_id: str, pipe: Optional[StableDiffusionPipeline] = None, @@ -155,7 +155,7 @@ def gen(self, inputs: dict) -> PIL.Image.Image: return image -class Im2ImPipe(BasePipe): +class Im2Im(BasePipe): def __init__(self, model_id, pipe: Optional[StableDiffusionImg2ImgPipeline] = None, **args): super().__init__(model_id=model_id, sd_pipe_class=StableDiffusionImg2ImgPipeline, pipe=pipe, **args) @@ -190,7 +190,7 @@ def gen(self, inputs: dict) -> PIL.Image.Image: return image -class Cond2ImPipe(BasePipe): +class Cond2Im(BasePipe): """ Base class for ControlNet pipelines """ @@ -225,7 +225,7 @@ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] = self.control_types = ctypes self._condition_image = None dtype = torch.float16 if 'torch_type' not in args else args['torch_type'] - cnets = [ControlNetModel.from_pretrained(Cond2ImPipe.cpath+Cond2ImPipe.cmodels[c], torch_dtype=dtype) for c in ctypes] + cnets = [ControlNetModel.from_pretrained(Cond2Im.cpath + Cond2Im.cmodels[c], torch_dtype=dtype) for c in ctypes] super().__init__(sd_pipe_class=StableDiffusionControlNetPipeline, model_id=model_id, pipe=pipe, controlnet=cnets, **args) # FIXME: do we need to setup this specific scheduler here? # should we pass its name in setup to super? @@ -264,14 +264,14 @@ def gen(self, inputs): return image -class ControlNet2ImPipe(Cond2ImPipe): +class ControlNet2Im(Cond2Im): """ ControlNet pipeline with conditional image generation """ def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] = None, ctypes=["soft"], **args): super().__init__(model_id=model_id, pipe=pipe, ctypes=ctypes, **args) - # The difference from Cond2ImPipe is that the conditional image is not + # The difference from Cond2Im is that the conditional image is not # taken as input but is obtained from an ordinary image, so this image # should be processed, and the processor depends on the conditioning type if "soft" in ctypes: @@ -333,13 +333,13 @@ def _proc_cimg(self, oriImg): # TODO: does it make sense to inherint it from Cond2Im or CIm2Im ? -class InpaintingPipe(BasePipe): +class Inpainting(BasePipe): def __init__(self, model_id, pipe: Optional[StableDiffusionControlNetPipeline] = None, **args): dtype = torch.float16 if 'torch_type' not in args else args['torch_type'] cnet = ControlNetModel.from_pretrained( - Cond2ImPipe.cpath+Cond2ImPipe.cmodels["inpaint"], torch_dtype=dtype) + Cond2Im.cpath + Cond2Im.cmodels["inpaint"], torch_dtype=dtype) super().__init__(sd_pipe_class=StableDiffusionControlNetInpaintPipeline, model_id=model_id, pipe=pipe, controlnet=cnet, **args) # FIXME: do we need to setup this specific scheduler here? # should we pass its name in setup to super? diff --git a/tests/pipe_test.py b/tests/pipe_test.py index 3346074..b349853 100644 --- a/tests/pipe_test.py +++ b/tests/pipe_test.py @@ -5,7 +5,7 @@ import torch import numpy -from multigen import Prompt2ImPipe, Cfgen, GenSession +from multigen import Prompt2Im, Cfgen, GenSession from dummy import DummyDiffusionPipeline @@ -22,7 +22,7 @@ def setUp(self): def test_basic_txt2im(self): model = "runwayml/stable-diffusion-v1-5" # create pipe - pipe = Prompt2ImPipe(model, pipe=self._pipeline) + pipe = Prompt2Im(model, pipe=self._pipeline) pipe.setup(width=512, height=512, guidance_scale=7, scheduler="DPMSolverMultistepScheduler") seed = 49045438434843 params = dict(prompt="a cube planet, cube-shaped, space photo, masterpiece", @@ -58,7 +58,7 @@ def test_with_session(self): ["8k RAW photo, masterpiece, super quality", "artwork", "unity 3D"], ["surrealism", "impressionism", "high tech", "cyberpunk"]] - pipe = Prompt2ImPipe(model, pipe=self._pipeline) + pipe = Prompt2Im(model, pipe=self._pipeline) pipe.setup(width=512, height=512, scheduler="DPMSolverMultistepScheduler") # remove directory if it exists dirname = "./gen_batch"