Upload extensions using SD-Hub extension

3dabe4a verified about 2 years ago

9.39 kB

import sys
from typing import Any, Callable, Union

from torch import nn
from torch.utils.hooks import RemovableHandle

from ldm.modules.diffusionmodules.openaimodel import (
    TimestepEmbedSequential,
)
from ldm.modules.attention import (
    SpatialTransformer,
    BasicTransformerBlock,
    CrossAttention,
    MemoryEfficientCrossAttention,
)
from ldm.modules.diffusionmodules.openaimodel import (
    ResBlock,
)
from modules.processing import StableDiffusionProcessing
from modules import shared

class ForwardHook:
    
    def __init__(self, module: nn.Module, fn: Callable[[nn.Module, Callable[..., Any], Any], Any]):
        self.o = module.forward
        self.fn = fn
        self.module = module
        self.module.forward = self.forward
    
    def remove(self):
        if self.module is not None and self.o is not None:
            self.module.forward = self.o
            self.module = None
            self.o = None
        self.fn = None
    
    def forward(self, *args, **kwargs):
        if self.module is not None and self.o is not None:
            if self.fn is not None:
                return self.fn(self.module, self.o, *args, **kwargs)
        return None
        

class SDHook:
    
    def __init__(self, enabled: bool):
        self._enabled = enabled
        self._handles: list[Union[RemovableHandle,ForwardHook]] = []
    
    @property
    def enabled(self):
        return self._enabled
    
    @property
    def batch_num(self):
        return shared.state.job_no
    
    @property
    def step_num(self):
        return shared.state.current_image_sampling_step
    
    def __enter__(self):
        if self.enabled:
            pass
    
    def __exit__(self, exc_type, exc_value, traceback):
        if self.enabled:
            for handle in self._handles:
                handle.remove()
            self._handles.clear()
            self.dispose()
    
    def dispose(self):
        pass
    
    def setup(
        self,
        p: StableDiffusionProcessing
    ):
        if not self.enabled:
            return
        
        wrapper = getattr(p.sd_model, "model", None)
        
        unet: Union[nn.Module,None] = getattr(wrapper, "diffusion_model", None) if wrapper is not None else None
        vae: Union[nn.Module,None] = getattr(p.sd_model, "first_stage_model", None)
        clip: Union[nn.Module,None] = getattr(p.sd_model, "cond_stage_model", None)
        
        assert unet is not None, "p.sd_model.diffusion_model is not found. broken model???"
        self._do_hook(p, p.sd_model, unet=unet, vae=vae, clip=clip) # type: ignore
        self.on_setup()
    
    def on_setup(self):
        pass
    
    def _do_hook(
        self,
        p: StableDiffusionProcessing,
        model: Any,
        unet: Union[nn.Module,None],
        vae: Union[nn.Module,None],
        clip: Union[nn.Module,None]
    ):
        assert model is not None, "empty model???"
        
        if clip is not None:
            self.hook_clip(p, clip)
        
        if unet is not None:
            self.hook_unet(p, unet)
        
        if vae is not None:
            self.hook_vae(p, vae)
    
    def hook_vae(
        self,
        p: StableDiffusionProcessing,
        vae: nn.Module
    ):
        pass

    def hook_unet(
        self,
        p: StableDiffusionProcessing,
        unet: nn.Module
    ):
        pass

    def hook_clip(
        self,
        p: StableDiffusionProcessing,
        clip: nn.Module
    ):
        pass

    def hook_layer(
        self,
        module: Union[nn.Module,Any],
        fn: Callable[[nn.Module, list, Any], Any]
    ):
        if not self.enabled:
            return
        
        assert module is not None
        assert isinstance(module, nn.Module)
        self._handles.append(module.register_forward_hook(fn))

    def hook_layer_pre(
        self,
        module: Union[nn.Module,Any],
        fn: Callable[[nn.Module, list], Any]
    ):
        if not self.enabled:
            return
        
        assert module is not None
        assert isinstance(module, nn.Module)
        self._handles.append(module.register_forward_pre_hook(fn))

    def hook_forward(
        self,
        module: Union[nn.Module,Any],
        fn: Callable[[nn.Module, Callable[..., Any], Any], Any]
    ):
        assert module is not None
        assert isinstance(module, nn.Module)
        self._handles.append(ForwardHook(module, fn))
    
    def log(self, msg: str):
        print(msg, file=sys.stderr)


# enumerate SpatialTransformer in TimestepEmbedSequential 
def each_transformer(unet_block: TimestepEmbedSequential):
    for block in unet_block.children():
        if isinstance(block, SpatialTransformer):
            yield block

# enumerate BasicTransformerBlock in SpatialTransformer
def each_basic_block(trans: SpatialTransformer):
    for block in trans.transformer_blocks.children():
        if isinstance(block, BasicTransformerBlock):
            yield block

# enumerate Attention Layers in TimestepEmbedSequential
# each_transformer + each_basic_block
def each_attns(unet_block: TimestepEmbedSequential):
    for n, trans in enumerate(each_transformer(unet_block)):
        for depth, basic_block in enumerate(each_basic_block(trans)):
            # attn1: Union[CrossAttention,MemoryEfficientCrossAttention]
            # attn2: Union[CrossAttention,MemoryEfficientCrossAttention]
            
            attn1, attn2 = basic_block.attn1, basic_block.attn2
            assert isinstance(attn1, CrossAttention) or isinstance(attn1, MemoryEfficientCrossAttention)
            assert isinstance(attn2, CrossAttention) or isinstance(attn2, MemoryEfficientCrossAttention)
            
            yield n, depth, attn1, attn2

def each_unet_attn_layers(unet: nn.Module):
    def get_attns(layer_index: int, block: TimestepEmbedSequential, format: str):
        for n, d, attn1, attn2 in each_attns(block):
            kwargs = {
                'layer_index': layer_index,
                'trans_index': n,
                'block_index': d
            }
            yield format.format(attn_name='sattn', **kwargs), attn1
            yield format.format(attn_name='xattn', **kwargs), attn2
    
    def enumerate_all(blocks: nn.ModuleList, format: str):
        for idx, block in enumerate(blocks.children()):
            if isinstance(block, TimestepEmbedSequential):
                yield from get_attns(idx, block, format)
    
    inputs: nn.ModuleList = unet.input_blocks           # type: ignore
    middle: TimestepEmbedSequential = unet.middle_block # type: ignore
    outputs: nn.ModuleList = unet.output_blocks         # type: ignore
    
    yield from enumerate_all(inputs, 'IN{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')
    yield from get_attns(0, middle, 'M{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')
    yield from enumerate_all(outputs, 'OUT{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')


def each_unet_transformers(unet: nn.Module):
    def get_trans(layer_index: int, block: TimestepEmbedSequential, format: str):
        for n, trans in enumerate(each_transformer(block)):
            kwargs = {
                'layer_index': layer_index,
                'block_index': n,
                'block_name': 'trans',
            }
            yield format.format(**kwargs), trans
    
    def enumerate_all(blocks: nn.ModuleList, format: str):
        for idx, block in enumerate(blocks.children()):
            if isinstance(block, TimestepEmbedSequential):
                yield from get_trans(idx, block, format)
    
    inputs: nn.ModuleList = unet.input_blocks           # type: ignore
    middle: TimestepEmbedSequential = unet.middle_block # type: ignore
    outputs: nn.ModuleList = unet.output_blocks         # type: ignore
    
    yield from enumerate_all(inputs, 'IN{layer_index:02}_{block_index:02}_{block_name}')
    yield from get_trans(0, middle, 'M{layer_index:02}_{block_index:02}_{block_name}')
    yield from enumerate_all(outputs, 'OUT{layer_index:02}_{block_index:02}_{block_name}')


def each_resblock(unet_block: TimestepEmbedSequential):
    for block in unet_block.children():
        if isinstance(block, ResBlock):
            yield block

def each_unet_resblock(unet: nn.Module):
    def get_resblock(layer_index: int, block: TimestepEmbedSequential, format: str):
        for n, res in enumerate(each_resblock(block)):
            kwargs = {
                'layer_index': layer_index,
                'block_index': n,
                'block_name': 'resblock',
            }
            yield format.format(**kwargs), res
    
    def enumerate_all(blocks: nn.ModuleList, format: str):
        for idx, block in enumerate(blocks.children()):
            if isinstance(block, TimestepEmbedSequential):
                yield from get_resblock(idx, block, format)
    
    inputs: nn.ModuleList = unet.input_blocks           # type: ignore
    middle: TimestepEmbedSequential = unet.middle_block # type: ignore
    outputs: nn.ModuleList = unet.output_blocks         # type: ignore
    
    yield from enumerate_all(inputs, 'IN{layer_index:02}_{block_index:02}_{block_name}')
    yield from get_resblock(0, middle, 'M{layer_index:02}_{block_index:02}_{block_name}')
    yield from enumerate_all(outputs, 'OUT{layer_index:02}_{block_index:02}_{block_name}')

	import sys
	from typing import Any, Callable, Union

	from torch import nn
	from torch.utils.hooks import RemovableHandle

	from ldm.modules.diffusionmodules.openaimodel import (
	TimestepEmbedSequential,
	)
	from ldm.modules.attention import (
	SpatialTransformer,
	BasicTransformerBlock,
	CrossAttention,
	MemoryEfficientCrossAttention,
	)
	from ldm.modules.diffusionmodules.openaimodel import (
	ResBlock,
	)
	from modules.processing import StableDiffusionProcessing
	from modules import shared

	class ForwardHook:

	def __init__(self, module: nn.Module, fn: Callable[[nn.Module, Callable[..., Any], Any], Any]):
	self.o = module.forward
	self.fn = fn
	self.module = module
	self.module.forward = self.forward

	def remove(self):
	if self.module is not None and self.o is not None:
	self.module.forward = self.o
	self.module = None
	self.o = None
	self.fn = None

	def forward(self, args, *kwargs):
	if self.module is not None and self.o is not None:
	if self.fn is not None:
	return self.fn(self.module, self.o, args, *kwargs)
	return None


	class SDHook:

	def __init__(self, enabled: bool):
	self._enabled = enabled
	self._handles: list[Union[RemovableHandle,ForwardHook]] = []

	@property
	def enabled(self):
	return self._enabled

	@property
	def batch_num(self):
	return shared.state.job_no

	@property
	def step_num(self):
	return shared.state.current_image_sampling_step

	def __enter__(self):
	if self.enabled:
	pass

	def __exit__(self, exc_type, exc_value, traceback):
	if self.enabled:
	for handle in self._handles:
	handle.remove()
	self._handles.clear()
	self.dispose()

	def dispose(self):
	pass

	def setup(
	self,
	p: StableDiffusionProcessing
	):
	if not self.enabled:
	return

	wrapper = getattr(p.sd_model, "model", None)

	unet: Union[nn.Module,None] = getattr(wrapper, "diffusion_model", None) if wrapper is not None else None
	vae: Union[nn.Module,None] = getattr(p.sd_model, "first_stage_model", None)
	clip: Union[nn.Module,None] = getattr(p.sd_model, "cond_stage_model", None)

	assert unet is not None, "p.sd_model.diffusion_model is not found. broken model???"
	self._do_hook(p, p.sd_model, unet=unet, vae=vae, clip=clip) # type: ignore
	self.on_setup()

	def on_setup(self):
	pass

	def _do_hook(
	self,
	p: StableDiffusionProcessing,
	model: Any,
	unet: Union[nn.Module,None],
	vae: Union[nn.Module,None],
	clip: Union[nn.Module,None]
	):
	assert model is not None, "empty model???"

	if clip is not None:
	self.hook_clip(p, clip)

	if unet is not None:
	self.hook_unet(p, unet)

	if vae is not None:
	self.hook_vae(p, vae)

	def hook_vae(
	self,
	p: StableDiffusionProcessing,
	vae: nn.Module
	):
	pass

	def hook_unet(
	self,
	p: StableDiffusionProcessing,
	unet: nn.Module
	):
	pass

	def hook_clip(
	self,
	p: StableDiffusionProcessing,
	clip: nn.Module
	):
	pass

	def hook_layer(
	self,
	module: Union[nn.Module,Any],
	fn: Callable[[nn.Module, list, Any], Any]
	):
	if not self.enabled:
	return

	assert module is not None
	assert isinstance(module, nn.Module)
	self._handles.append(module.register_forward_hook(fn))

	def hook_layer_pre(
	self,
	module: Union[nn.Module,Any],
	fn: Callable[[nn.Module, list], Any]
	):
	if not self.enabled:
	return

	assert module is not None
	assert isinstance(module, nn.Module)
	self._handles.append(module.register_forward_pre_hook(fn))

	def hook_forward(
	self,
	module: Union[nn.Module,Any],
	fn: Callable[[nn.Module, Callable[..., Any], Any], Any]
	):
	assert module is not None
	assert isinstance(module, nn.Module)
	self._handles.append(ForwardHook(module, fn))

	def log(self, msg: str):
	print(msg, file=sys.stderr)


	# enumerate SpatialTransformer in TimestepEmbedSequential
	def each_transformer(unet_block: TimestepEmbedSequential):
	for block in unet_block.children():
	if isinstance(block, SpatialTransformer):
	yield block

	# enumerate BasicTransformerBlock in SpatialTransformer
	def each_basic_block(trans: SpatialTransformer):
	for block in trans.transformer_blocks.children():
	if isinstance(block, BasicTransformerBlock):
	yield block

	# enumerate Attention Layers in TimestepEmbedSequential
	# each_transformer + each_basic_block
	def each_attns(unet_block: TimestepEmbedSequential):
	for n, trans in enumerate(each_transformer(unet_block)):
	for depth, basic_block in enumerate(each_basic_block(trans)):
	# attn1: Union[CrossAttention,MemoryEfficientCrossAttention]
	# attn2: Union[CrossAttention,MemoryEfficientCrossAttention]

	attn1, attn2 = basic_block.attn1, basic_block.attn2
	assert isinstance(attn1, CrossAttention) or isinstance(attn1, MemoryEfficientCrossAttention)
	assert isinstance(attn2, CrossAttention) or isinstance(attn2, MemoryEfficientCrossAttention)

	yield n, depth, attn1, attn2

	def each_unet_attn_layers(unet: nn.Module):
	def get_attns(layer_index: int, block: TimestepEmbedSequential, format: str):
	for n, d, attn1, attn2 in each_attns(block):
	kwargs = {
	'layer_index': layer_index,
	'trans_index': n,
	'block_index': d
	}
	yield format.format(attn_name='sattn', **kwargs), attn1
	yield format.format(attn_name='xattn', **kwargs), attn2

	def enumerate_all(blocks: nn.ModuleList, format: str):
	for idx, block in enumerate(blocks.children()):
	if isinstance(block, TimestepEmbedSequential):
	yield from get_attns(idx, block, format)

	inputs: nn.ModuleList = unet.input_blocks # type: ignore
	middle: TimestepEmbedSequential = unet.middle_block # type: ignore
	outputs: nn.ModuleList = unet.output_blocks # type: ignore

	yield from enumerate_all(inputs, 'IN{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')
	yield from get_attns(0, middle, 'M{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')
	yield from enumerate_all(outputs, 'OUT{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')


	def each_unet_transformers(unet: nn.Module):
	def get_trans(layer_index: int, block: TimestepEmbedSequential, format: str):
	for n, trans in enumerate(each_transformer(block)):
	kwargs = {
	'layer_index': layer_index,
	'block_index': n,
	'block_name': 'trans',
	}
	yield format.format(**kwargs), trans

	def enumerate_all(blocks: nn.ModuleList, format: str):
	for idx, block in enumerate(blocks.children()):
	if isinstance(block, TimestepEmbedSequential):
	yield from get_trans(idx, block, format)

	inputs: nn.ModuleList = unet.input_blocks # type: ignore
	middle: TimestepEmbedSequential = unet.middle_block # type: ignore
	outputs: nn.ModuleList = unet.output_blocks # type: ignore

	yield from enumerate_all(inputs, 'IN{layer_index:02}_{block_index:02}_{block_name}')
	yield from get_trans(0, middle, 'M{layer_index:02}_{block_index:02}_{block_name}')
	yield from enumerate_all(outputs, 'OUT{layer_index:02}_{block_index:02}_{block_name}')


	def each_resblock(unet_block: TimestepEmbedSequential):
	for block in unet_block.children():
	if isinstance(block, ResBlock):
	yield block

	def each_unet_resblock(unet: nn.Module):
	def get_resblock(layer_index: int, block: TimestepEmbedSequential, format: str):
	for n, res in enumerate(each_resblock(block)):
	kwargs = {
	'layer_index': layer_index,
	'block_index': n,
	'block_name': 'resblock',
	}
	yield format.format(**kwargs), res

	def enumerate_all(blocks: nn.ModuleList, format: str):
	for idx, block in enumerate(blocks.children()):
	if isinstance(block, TimestepEmbedSequential):
	yield from get_resblock(idx, block, format)

	inputs: nn.ModuleList = unet.input_blocks # type: ignore
	middle: TimestepEmbedSequential = unet.middle_block # type: ignore
	outputs: nn.ModuleList = unet.output_blocks # type: ignore

	yield from enumerate_all(inputs, 'IN{layer_index:02}_{block_index:02}_{block_name}')
	yield from get_resblock(0, middle, 'M{layer_index:02}_{block_index:02}_{block_name}')
	yield from enumerate_all(outputs, 'OUT{layer_index:02}_{block_index:02}_{block_name}')