InvokeAI 2.2.4.5__py3-none-any.whl

ldm/invoke/generator/txt2img2img.py

@@ -0,0 +1,182 @@
+'''
+ldm.invoke.generator.txt2img inherits from ldm.invoke.generator
+'''
+
+import torch
+import numpy as  np
+import math
+import gc
+from ldm.invoke.generator.base import Generator
+from ldm.models.diffusion.ddim import DDIMSampler
+from ldm.invoke.generator.omnibus import Omnibus
+from ldm.models.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent
+from PIL import Image
+
+class Txt2Img2Img(Generator):
+    def __init__(self, model, precision):
+        super().__init__(model, precision)
+        self.init_latent = None    # for get_noise()
+
+    @torch.no_grad()
+    def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta,
+                       conditioning,width,height,strength,step_callback=None,**kwargs):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it
+        kwargs are 'width' and 'height'
+        """
+        uc, c, extra_conditioning_info = conditioning
+        scale_dim = min(width, height)
+        scale = 512 / scale_dim
+
+        init_width = math.ceil(scale * width / 64) * 64
+        init_height = math.ceil(scale * height / 64) * 64
+
+        @torch.no_grad()
+        def make_image(x_T):
+
+            shape = [
+                self.latent_channels,
+                init_height // self.downsampling_factor,
+                init_width // self.downsampling_factor,
+            ]
+
+            sampler.make_schedule(
+                    ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
+            )
+
+            #x = self.get_noise(init_width, init_height)
+            x = x_T
+
+            if self.free_gpu_mem and self.model.model.device != self.model.device:
+                self.model.model.to(self.model.device)
+
+            samples, _ = sampler.sample(
+                batch_size                   = 1,
+                S                            = steps,
+                x_T                          = x,
+                conditioning                 = c,
+                shape                        = shape,
+                verbose                      = False,
+                unconditional_guidance_scale = cfg_scale,
+                unconditional_conditioning   = uc,
+                eta                          = ddim_eta,
+                img_callback                 = step_callback,
+                extra_conditioning_info      = extra_conditioning_info
+            )
+
+            print(
+                  f"\n>> Interpolating from {init_width}x{init_height} to {width}x{height} using DDIM sampling"
+                 )
+
+            # resizing
+            samples = torch.nn.functional.interpolate(
+                samples,
+                size=(height // self.downsampling_factor, width // self.downsampling_factor),
+                mode="bilinear"
+            )
+
+            t_enc = int(strength * steps)
+            ddim_sampler = DDIMSampler(self.model, device=self.model.device)
+            ddim_sampler.make_schedule(
+                    ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
+            )
+
+            z_enc = ddim_sampler.stochastic_encode(
+                samples,
+                torch.tensor([t_enc-1]).to(self.model.device),
+                noise=self.get_noise(width,height,False)
+            )
+
+            # decode it
+            samples = ddim_sampler.decode(
+                z_enc,
+                c,
+                t_enc,
+                img_callback = step_callback,
+                unconditional_guidance_scale=cfg_scale,
+                unconditional_conditioning=uc,
+                extra_conditioning_info=extra_conditioning_info,
+                all_timesteps_count=steps
+            )
+
+            if self.free_gpu_mem:
+                self.model.model.to('cpu')
+                self.model.cond_stage_model.device = 'cpu'
+                self.model.cond_stage_model.to('cpu')
+                gc.collect()
+                torch.cuda.empty_cache()
+
+            return self.sample_to_image(samples)
+
+        # in the case of the inpainting model being loaded, the trick of
+        # providing an interpolated latent doesn't work, so we transiently
+        # create a 512x512 PIL image, upscale it, and run the inpainting
+        # over it in img2img mode. Because the inpaing model is so conservative
+        # it doesn't change the image (much)
+        def inpaint_make_image(x_T):
+            omnibus = Omnibus(self.model,self.precision)
+            result = omnibus.generate(
+                prompt,
+                sampler=sampler,
+                width=init_width,
+                height=init_height,
+                step_callback=step_callback,
+                steps = steps,
+                cfg_scale = cfg_scale,
+                ddim_eta = ddim_eta,
+                conditioning = conditioning,
+                **kwargs
+            )
+            assert result is not None and len(result)>0,'** txt2img failed **'
+            image = result[0][0]
+            interpolated_image = image.resize((width,height),resample=Image.Resampling.LANCZOS)
+            print(kwargs.pop('init_image',None))
+            result = omnibus.generate(
+                prompt,
+                sampler=sampler,
+                init_image=interpolated_image,
+                width=width,
+                height=height,
+                seed=result[0][1],
+                step_callback=step_callback,
+                steps = steps,
+                cfg_scale = cfg_scale,
+                ddim_eta = ddim_eta,
+                conditioning = conditioning,
+                **kwargs
+                )
+            return result[0][0]
+
+        if sampler.uses_inpainting_model():
+            return inpaint_make_image
+        else:
+            return make_image
+
+    # returns a tensor filled with random numbers from a normal distribution
+    def get_noise(self,width,height,scale = True):
+        # print(f"Get noise: {width}x{height}")
+        if scale:
+            trained_square = 512 * 512
+            actual_square = width * height
+            scale = math.sqrt(trained_square / actual_square)
+            scaled_width = math.ceil(scale * width / 64) * 64
+            scaled_height = math.ceil(scale * height / 64) * 64
+        else:
+            scaled_width = width
+            scaled_height = height
+
+        device      = self.model.device
+        if self.use_mps_noise or device.type == 'mps':
+            return torch.randn([1,
+                                self.latent_channels,
+                                scaled_height // self.downsampling_factor,
+                                scaled_width  // self.downsampling_factor],
+                                device='cpu').to(device)
+        else:
+            return torch.randn([1,
+                                self.latent_channels,
+                                scaled_height // self.downsampling_factor,
+                                scaled_width  // self.downsampling_factor],
+                                device=device)
+

ldm/invoke/globals.py

@@ -0,0 +1,31 @@
+'''
+ldm.invoke.globals defines a small number of global variables that would
+otherwise have to be passed through long and complex call chains.
+
+It defines a Namespace object named "Globals" that contains
+the attributes:
+
+  - root           - the root directory under which "models" and "outputs" can be found
+  - initfile       - path to the initialization file
+  - try_patchmatch - option to globally disable loading of 'patchmatch' module
+'''
+
+import os
+import os.path as osp
+from argparse import Namespace
+
+Globals = Namespace()
+
+# This is usually overwritten by the command line and/or environment variables
+if os.environ.get('INVOKEAI_ROOT'):
+    Globals.root = osp.abspath(os.environ.get('INVOKEAI_ROOT'))
+elif os.environ.get('VIRTUAL_ENV'):
+    Globals.root = osp.abspath(osp.join(os.environ.get('VIRTUAL_ENV'), '..'))
+else:
+    Globals.root = osp.abspath(osp.expanduser('~/invokeai'))
+
+# Where to look for the initialization file
+Globals.initfile = 'invokeai.init'
+
+# Try loading patchmatch
+Globals.try_patchmatch = True

ldm/invoke/image_util.py

@@ -0,0 +1,70 @@
+from math import sqrt, floor, ceil
+from PIL import Image
+
+class InitImageResizer():
+    """Simple class to create resized copies of an Image while preserving the aspect ratio."""
+    def __init__(self,Image):
+        self.image = Image
+
+    def resize(self,width=None,height=None) -> Image:
+        """
+        Return a copy of the image resized to fit within
+        a box width x height. The aspect ratio is
+        maintained. If neither width nor height are provided,
+        then returns a copy of the original image. If one or the other is
+        provided, then the other will be calculated from the
+        aspect ratio.
+
+        Everything is floored to the nearest multiple of 64 so
+        that it can be passed to img2img()
+        """
+        im    = self.image
+
+        ar = im.width/float(im.height)
+
+        # Infer missing values from aspect ratio
+        if not(width or height): # both missing
+            width  = im.width
+            height = im.height
+        elif not height:           # height missing
+            height = int(width/ar)
+        elif not width:            # width missing
+            width  = int(height*ar)
+
+        w_scale = width/im.width
+        h_scale = height/im.height
+        scale = min(w_scale,h_scale)
+        (rw,rh) = (int(scale*im.width),int(scale*im.height))
+
+        #round everything to multiples of 64
+        width,height,rw,rh = map(
+            lambda x: x-x%64, (width,height,rw,rh)
+        )
+
+        # no resize necessary, but return a copy
+        if im.width == width and im.height == height:
+            return im.copy()
+
+        # otherwise resize the original image so that it fits inside the bounding box
+        resized_image = self.image.resize((rw,rh),resample=Image.Resampling.LANCZOS)
+        return resized_image
+
+def make_grid(image_list, rows=None, cols=None):
+    image_cnt = len(image_list)
+    if None in (rows, cols):
+        rows = floor(sqrt(image_cnt))  # try to make it square
+        cols = ceil(image_cnt / rows)
+    width = image_list[0].width
+    height = image_list[0].height
+
+    grid_img = Image.new('RGB', (width * cols, height * rows))
+    i = 0
+    for r in range(0, rows):
+        for c in range(0, cols):
+            if i >= len(image_list):
+                break
+            grid_img.paste(image_list[i], (c * width, r * height))
+            i = i + 1
+
+    return grid_img
+

ldm/invoke/log.py

@@ -0,0 +1,66 @@
+"""
+Functions for better format logging
+    write_log -- logs the name of the output image, prompt, and prompt args to the terminal and different types of file
+        1 write_log_message -- Writes a message to the console
+        2 write_log_files -- Writes a message to files
+        2.1 write_log_default -- File in plain text
+        2.2 write_log_txt -- File in txt format
+        2.3 write_log_markdown -- File in markdown format
+"""
+
+import os
+
+
+def write_log(results, log_path, file_types, output_cntr):
+    """
+    logs the name of the output image, prompt, and prompt args to the terminal and files
+    """
+    output_cntr = write_log_message(results, output_cntr)
+    write_log_files(results, log_path, file_types)
+    return output_cntr
+
+
+def write_log_message(results, output_cntr):
+    """logs to the terminal"""
+    if len(results) == 0:
+        return output_cntr
+    log_lines = [f"{path}: {prompt}\n" for path, prompt in results]
+    if len(log_lines)>1:
+        subcntr = 1
+        for l in log_lines:
+           print(f"[{output_cntr}.{subcntr}] {l}", end="")
+           subcntr += 1
+    else:
+           print(f"[{output_cntr}] {log_lines[0]}", end="")
+    return output_cntr+1
+
+def write_log_files(results, log_path, file_types):
+    for file_type in file_types:
+        if file_type == "txt":
+            write_log_txt(log_path, results)
+        elif file_type == "md" or file_type == "markdown":
+            write_log_markdown(log_path, results)
+        else:
+            print(f"'{file_type}' format is not supported, so write in plain text")
+            write_log_default(log_path, results, file_type)
+
+
+def write_log_default(log_path, results, file_type):
+    plain_txt_lines = [f"{path}: {prompt}\n" for path, prompt in results]
+    with open(log_path + "." + file_type, "a", encoding="utf-8") as file:
+        file.writelines(plain_txt_lines)
+
+
+def write_log_txt(log_path, results):
+    txt_lines = [f"{path}: {prompt}\n" for path, prompt in results]
+    with open(log_path + ".txt", "a", encoding="utf-8") as file:
+        file.writelines(txt_lines)
+
+
+def write_log_markdown(log_path, results):
+    md_lines = []
+    for path, prompt in results:
+        file_name = os.path.basename(path)
+        md_lines.append(f"## {file_name}\n![]({file_name})\n\n{prompt}\n")
+    with open(log_path + ".md", "a", encoding="utf-8") as file:
+        file.writelines(md_lines)