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 cc5aeeb..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 = CIm2ImPipe(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 c580d70..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,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 = 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 5ec8e51..c3648d4 100755 --- a/multigen/pipes.py +++ b/multigen/pipes.py @@ -1,4 +1,6 @@ import importlib + +import PIL import torch from PIL import Image @@ -124,7 +126,8 @@ def setup(self, steps=50, **args): # TODO? add scheduler to config? self.try_set_scheduler(dict(scheduler=args['scheduler'])) -class Prompt2ImPipe(BasePipe): + +class Prompt2Im(BasePipe): def __init__(self, model_id: str, pipe: Optional[StableDiffusionPipeline] = None, @@ -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) @@ -152,24 +155,23 @@ def gen(self, inputs): 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) 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) @@ -188,10 +190,14 @@ def gen(self, inputs): return image -class Cond2ImPipe(BasePipe): +class Cond2Im(BasePipe): + """ + Base class for ControlNet pipelines + """ # TODO: set path cpath = "./models-cn/" + cmodels = { "canny": "sd-controlnet-canny", "pose": "control_v11p_sd15_openpose", @@ -201,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, @@ -215,23 +222,23 @@ 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(CIm2ImPipe.cpath+CIm2ImPipe.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? 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 = [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, @@ -244,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 @@ -257,12 +264,14 @@ def gen(self, inputs): return image -class CIm2ImPipe(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: @@ -280,15 +289,15 @@ 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])) 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)] @@ -324,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/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 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"