File size: 2,515 Bytes
c10aebf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from PIL import Image
from sd_bmab.base import cache
from sd_bmab.base import ProcessorBase
from sd_bmab.base.context import Context
from sd_bmab.external import load_external_module
class ICLight(ProcessorBase):
def __init__(self) -> None:
super().__init__()
self.iclight_opt = {}
self.enabled = False
self.enable_before_upscale = False
self.style = 'normal'
self.prompt = ''
self.preference = 'None'
self.use_background_image = False
self.blending = 0.5
def preprocess(self, context: Context, image: Image):
self.iclight_opt = context.args.get('module_config', {}).get('iclight', {})
self.enabled = self.iclight_opt.get('enabled', self.enabled)
self.enable_before_upscale = self.iclight_opt.get('enable_before_upscale', self.enable_before_upscale)
self.style = self.iclight_opt.get('style', self.style)
self.prompt = self.iclight_opt.get('prompt', self.prompt)
self.preference = self.iclight_opt.get('preference', self.preference)
self.use_background_image = self.iclight_opt.get('use_background_image', self.use_background_image)
self.blending = self.iclight_opt.get('blending', self.blending)
return self.enabled
def process(self, context: Context, image: Image):
mod = load_external_module('iclight', 'bmabiclight')
bg_image = self.get_background_image() if self.use_background_image else None
return mod.bmab_relight(context, self.style, image, bg_image, self.prompt, self.blending, self.preference)
def postprocess(self, context: Context, image: Image):
pass
@staticmethod
def get_background_image():
img = cache.get_image_from_cache('iclight_background.png')
if img is not None:
return img
return Image.new('RGB', (512, 768), (0, 0, 0))
@staticmethod
def put_backgound_image(img):
cache.put_image_to_cache('iclight_background.png', img)
@staticmethod
def get_styles():
return ['intensive', 'less intensive', 'normal', 'soft']
class ICLightBeforeUpsacle(ICLight):
def preprocess(self, context: Context, image: Image):
super().preprocess(context, image)
return self.enabled and self.enable_before_upscale
class ICLightAfterUpsacle(ICLight):
def preprocess(self, context: Context, image: Image):
super().preprocess(context, image)
return self.enabled and not self.enable_before_upscale
|