feature: tile mode can now be specified per-prompt

README.md

@@ -117,7 +117,7 @@ from imaginairy import imagine, imagine_image_files, ImaginePrompt, WeightedProm
 
 url = "https://upload.wikimedia.org/wikipedia/commons/thumb/6/6c/Thomas_Cole_-_Architect%E2%80%99s_Dream_-_Google_Art_Project.jpg/540px-Thomas_Cole_-_Architect%E2%80%99s_Dream_-_Google_Art_Project.jpg"
 prompts = [
-    ImaginePrompt("a scenic landscape", seed=1),
+    ImaginePrompt("a scenic landscape", seed=1, upscale=True),
     ImaginePrompt("a bowl of fruit"),
     ImaginePrompt([
         WeightedPrompt("cat", weight=1),
@@ -133,7 +133,8 @@ prompts = [
         mask_prompt="fruit|stems",
         mask_mode="replace",
         mask_expansion=3
-    )
+    ),
+    ImaginePrompt("strawberries", tile_mode=True),
 ]
 for result in imagine(prompts):
     # do something
@@ -162,7 +163,11 @@ docker build . -t imaginairy
 docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface -v $HOME/.cache/torch:/root/.cache/torch -v `pwd`/outputs:/outputs imaginairy /bin/bash
 ```
 
+## Running on Google Colab
+[Example Colab](https://colab.research.google.com/drive/1rOvQNs0Cmn_yU1bKWjCOHzGVDgZkaTtO?usp=sharing)
+
 ## ChangeLog
+ - tile mode can now be specified per-prompt
 
 **1.5.3**
  - fix: missing config file for describe feature

imaginairy/api.py

@@ -73,31 +73,11 @@ def load_model_from_config(config):
     return model
 
 
-def patch_conv(**patch):
-    """
-    Patch to enable tiling mode
-
-    https://github.com/replicate/cog-stable-diffusion/compare/main...TomMoore515:material_stable_diffusion:main
-    """
-    cls = torch.nn.Conv2d
-    init = cls.__init__
-
-    def __init__(self, *args, **kwargs):
-        return init(self, *args, **kwargs, **patch)
-
-    cls.__init__ = __init__
-
-
 @lru_cache()
-def load_model(tile_mode=False):
-    if tile_mode:
-        # generated images are tileable
-        patch_conv(padding_mode="circular")
-
+def load_model():
     config = "configs/stable-diffusion-v1.yaml"
     config = OmegaConf.load(f"{LIB_PATH}/{config}")
     model = load_model_from_config(config)
-
     model = model.to(get_device())
     return model
 
@@ -111,7 +91,6 @@ def imagine_image_files(
     ddim_eta=0.0,
     record_step_images=False,
     output_file_extension="jpg",
-    tile_mode=False,
     print_caption=False,
 ):
     big_path = os.path.join(outdir, "upscaled")
@@ -139,7 +118,6 @@ def imagine_image_files(
         precision=precision,
         ddim_eta=ddim_eta,
         img_callback=_record_step if record_step_images else None,
-        tile_mode=tile_mode,
         add_caption=print_caption,
     ):
         prompt = result.prompt
@@ -164,11 +142,10 @@ def imagine(
     precision="autocast",
     ddim_eta=0.0,
     img_callback=None,
-    tile_mode=False,
     half_mode=None,
     add_caption=False,
 ):
-    model = load_model(tile_mode=tile_mode)
+    model = load_model()
 
     # only run half-mode on cuda. run it by default
     half_mode = half_mode is None and get_device() == "cuda"
@@ -194,6 +171,7 @@ def imagine(
             ):
                 logger.info(f"Generating {prompt.prompt_description()}")
                 seed_everything(prompt.seed)
+                model.tile_mode(prompt.tile_mode)
 
                 uc = None
                 if prompt.prompt_strength != 1.0:

imaginairy/cmds.py

@@ -121,7 +121,7 @@ def configure_logging(level="INFO"):
 @click.option(
     "--tile",
     is_flag=True,
-    help="Any images rendered will be tileable.  Unfortunately cannot be controlled at the per-image level yet",
+    help="Any images rendered will be tileable.",
 )
 @click.option(
     "--mask-image",
@@ -190,7 +190,7 @@ def imagine_cmd(
         init_image = LazyLoadingImage(url=init_image)
 
     prompts = []
-    load_model(tile_mode=tile)
+    load_model()
     for _ in range(repeats):
         for prompt_text in prompt_texts:
             prompt = ImaginePrompt(
@@ -209,6 +209,7 @@ def imagine_cmd(
                 mask_mode=mask_mode,
                 upscale=upscale,
                 fix_faces=fix_faces,
+                tile_mode=tile,
             )
             prompts.append(prompt)
 
@@ -217,7 +218,6 @@ def imagine_cmd(
         outdir=outdir,
         ddim_eta=ddim_eta,
         record_step_images="images" in show_work,
-        tile_mode=tile,
         output_file_extension="png",
         print_caption=caption,
     )

imaginairy/modules/diffusion/ddpm.py

@@ -273,6 +273,18 @@ class LatentDiffusion(DDPM):
             self.init_from_ckpt(ckpt_path, ignore_keys)
             self.restarted_from_ckpt = True
 
+        # store initial padding mode so we can switch to 'circular'
+        # when we want tiled images
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                m._initial_padding_mode = m.padding_mode
+
+    def tile_mode(self, enabled):
+        """For creating seamless tiles"""
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                m.padding_mode = "circular" if enabled else m._initial_padding_mode
+
     def make_cond_schedule(
         self,
     ):

imaginairy/schema.py

@@ -103,6 +103,7 @@ class ImaginePrompt:
         fix_faces=False,
         sampler_type="PLMS",
         conditioning=None,
+        tile_mode=False,
     ):
         prompt = prompt if prompt is not None else "a scenic landscape"
         if isinstance(prompt, str):
@@ -131,6 +132,7 @@ class ImaginePrompt:
         self.mask_image = mask_image
         self.mask_mode = mask_mode
         self.mask_expansion = mask_expansion
+        self.tile_mode = tile_mode
 
     @property
     def prompt_text(self):