refactoring and typing

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)

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)

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)

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)

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()
 

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

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,33 +146,32 @@ 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)
         image = self.pipe(**inputs).images[0]
         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,19 +181,23 @@ 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)
         image = self.pipe(**inputs).images[0]
         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?

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
 

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"