Merge pull request #48 from noskill/imim2

support im2im pipeline in worker

multigen/loader.py

@@ -1,4 +1,5 @@
 from typing import Type
+import torch
 import logging
 from diffusers import DiffusionPipeline, StableDiffusionControlNetPipeline, StableDiffusionXLControlNetPipeline
 
@@ -12,7 +13,7 @@ class for loading diffusion pipelines from files.
     def __init__(self):
         self._pipes = dict()
 
-    def load_pipeline(self, cls: Type[DiffusionPipeline], path, **additional_args):
+    def load_pipeline(self, cls: Type[DiffusionPipeline], path, torch_dtype=torch.float16, **additional_args):
         for key, pipe in self._pipes.items():
             if key == path:
                 components = pipe.components
@@ -25,7 +26,7 @@ def load_pipeline(self, cls: Type[DiffusionPipeline], path, **additional_args):
                 # but we don't need it
                 if 'controlnet' in components:
                     components.pop('controlnet')
-                return cls(**components, **additional_args)
+                return cls(**components, **additional_args).to(torch_dtype)
 
 
         if path.endswith('safetensors'):

multigen/worker.py

@@ -56,7 +56,7 @@ def _update(sess, job, gs):
                 if 'MaskedIm2ImPipe' in class_name:
                     pipe.setup(**data, original_image=str(images[0]),
                             image_painted=str(images[1]))
-                elif 'Cond2ImPipe' in class_name:
+                elif any([x in class_name for x in ('Im2ImPipe', 'Cond2ImPipe')]):
                     pipe.setup(**data, fimage=str(images[0]))
                 else:
                     pipe.setup(**data)

multigen/worker_base.py

@@ -7,7 +7,7 @@
 import yaml
 from pathlib import Path
 import logging
-from .pipes import Prompt2ImPipe, MaskedIm2ImPipe, Cond2ImPipe
+from .pipes import Prompt2ImPipe, MaskedIm2ImPipe, Cond2ImPipe, Im2ImPipe
 from .loader import Loader