Upload folder using huggingface_hub

gray-inpaint/config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "SDGrayInpaintModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_sd_gray_inpaint.SDGrayInpaintConfig",
+    "AutoModel": "modeling_sd_gray_inpaint.SDGrayInpaintModel"
+  },
+  "base_model": "stabilityai/stable-diffusion-2-inpainting",
+  "height": 512,
+  "model_type": "sd_gray_inpaint",
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "width": 512
+}

gray-inpaint/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c6d964dca7f33a3a87e90056e8ab617efeabd99e3cfcea71f73d459b133f231
+size 4055354432

gray-inpaint/modeling_sd_gray_inpaint.py

@@ -0,0 +1,98 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from copy import deepcopy
+from torchvision.transforms.functional import rgb_to_grayscale
+import segmentation_models_pytorch as smp
+from diffusers import StableDiffusionInpaintPipeline
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import PretrainedConfig, PreTrainedModel
+
+class SDGrayInpaintConfig(PretrainedConfig):
+    model_type = "sd_gray_inpaint"
+    def __init__(
+        self,
+        base_model="stabilityai/stable-diffusion-2-inpainting",
+        height=512,
+        width=512,
+        **kwargs
+    ):
+        self.base_model=base_model
+        self.height=height
+        self.width=width
+        super().__init__(**kwargs)
+
+class SDGrayInpaintModel(PreTrainedModel):
+    config_class = SDGrayInpaintConfig
+    def __init__(self, config):
+        super().__init__(config)
+        pipe = StableDiffusionInpaintPipeline.from_pretrained(config.base_model)
+        self.mask_predictor = smp.Unet(
+            encoder_name="mit_b4",        
+            encoder_weights="imagenet",     
+            in_channels=3,                  
+            classes=1,                      
+        )
+        self.image_processor = pipe.image_processor
+        self.scheduler = pipe.scheduler
+        self.unet = pipe.unet
+        self.vae = pipe.vae
+        self.prompt_embeds = nn.Parameter(torch.randn(1,77,1024))
+        self.height=config.height
+        self.width=config.width
+
+    def forward(
+        self, 
+        images_gray_masked, 
+        masks=None, 
+        num_inference_steps=250,
+        seed=42,
+        input_type='pil', 
+        output_type='pil'
+    ):
+        generator = torch.Generator()
+        generator.manual_seed(seed)
+        if input_type=='pil':
+            images_gray_masked = self.image_processor.preprocess(images_gray_masked, height=self.height, width=self.width).float()
+        elif input_type=='pt':
+            images_gray_masked=images_gray_masked
+        else:
+            raise ValueError('unsupported input_type')
+        images_gray_masked = images_gray_masked.to(self.vae.device)
+        if masks is None:
+            masks_logits = self.mask_predictor(images_gray_masked)
+            masks = (torch.sigmoid(masks_logits)>0.5)*1.
+        masks = masks.float().to(self.vae.device)
+        B, C, H, W = images_gray_masked.shape
+        prompt_embeds = self.prompt_embeds.repeat(B,1,1)
+
+        scheduler = deepcopy(self.scheduler)
+        scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=self.vae.device)
+        masked_image_latents = self.vae.encode(images_gray_masked).latent_dist.mode() * self.vae.config.scaling_factor
+        mask_latents = F.interpolate(masks, size=(self.unet.config.sample_size, self.unet.config.sample_size))
+        latents = randn_tensor(masked_image_latents.shape, generator=generator).to(self.device) * self.scheduler.init_noise_sigma
+        for t in scheduler.timesteps:
+            latents = scheduler.scale_model_input(latents, t)
+            latent_model_input = torch.cat([latents, mask_latents, masked_image_latents], dim=1)
+            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds)[0]
+            latents = scheduler.step(noise_pred, t, latents)[0]
+        latents = latents / self.vae.config.scaling_factor
+        images_gray_restored = self.vae.decode(latents.detach())[0]
+        images_gray_restored = images_gray_masked * (1-masks) + images_gray_restored.detach() * masks
+        images_gray_restored = rgb_to_grayscale(images_gray_restored)
+        
+        if output_type=='pil':
+            images_gray_restored = self.image_processor.postprocess(images_gray_restored)
+        elif output_type=='np':
+            images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'np')
+        elif output_type=='pt':
+            images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'pt')
+        elif output_type=='none':
+            images_gray_restored = images_gray_restored
+        else:
+            raise ValueError('unsupported output_type')
+
+        return images_gray_restored
+
+

gray2rgb/config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "SeResVaeModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_seresvae.SeResVaeConfig",
+    "AutoModel": "modeling_seresvae.SeResVaeModel"
+  },
+  "base_model": "stabilityai/stable-diffusion-2-1",
+  "height": 512,
+  "model_type": "seresvae",
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "width": 512
+}

gray2rgb/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:657257a0969eec19b5a3ff2629500454e1941615ae646fa72f2d88e9d41be737
+size 3799014812

gray2rgb/modeling_seresvae.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+
+from diffusers import AutoencoderKL, UNet2DConditionModel
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from transformers import PretrainedConfig, PreTrainedModel
+
+class SEPath(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction=16):
+        super(SEPath, self).__init__()
+        self.fc = nn.Sequential(
+            nn.Linear(in_channels, in_channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_channels // reduction, out_channels, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, in_tensor, out_tensor):
+        B, C, H, W = in_tensor.size()
+        # Squeeze operation
+        x = in_tensor.view(B, C, -1).mean(dim=2)
+        # Excitation operation
+        x = self.fc(x).unsqueeze(2).unsqueeze(2)
+
+        return out_tensor * x
+
+class SeResVaeConfig(PretrainedConfig):
+    model_type = "seresvae"
+    def __init__(
+        self,
+        base_model="stabilityai/stable-diffusion-2-1",
+        height=512,
+        width=512,
+        **kwargs
+    ):
+        self.base_model=base_model
+        self.height=height
+        self.width=width
+        super().__init__(**kwargs)
+
+class SeResVaeModel(PreTrainedModel):
+    config_class = SeResVaeConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.image_processor = VaeImageProcessor()
+        self.vae = AutoencoderKL.from_pretrained(config.base_model, subfolder='vae')
+        self.unet = UNet2DConditionModel.from_pretrained(config.base_model, subfolder='unet')
+        self.se_paths = nn.ModuleList([SEPath(8,4), SEPath(512,512), SEPath(512,512), SEPath(256,512), SEPath(128,256)])
+        self.prompt_embeds = nn.Parameter(torch.randn(1,77,1024))
+        self.height=config.height
+        self.width=config.width
+
+    def forward(self, images_gray, input_type='pil', output_type='pil'):
+        if input_type=='pil':
+            images_gray = self.image_processor.preprocess(images_gray, height=self.height, width=self.width).float()
+        elif input_type=='pt':
+            images_gray=images_gray
+        else:
+            raise ValueError('unsupported input_type')
+        images_gray = images_gray.to(self.vae.device)
+        B, C, H, W = images_gray.shape
+        prompt_embeds = self.prompt_embeds.repeat(B,1,1)
+
+        posterior, encode_residual = self.encode_with_residual(images_gray)
+        latents = posterior.mode()
+        t = torch.LongTensor([500]).repeat(B).to(self.vae.device)
+        noise_pred = self.unet(latents, t, encoder_hidden_states=prompt_embeds)[0]
+        denoised_latents = latents - noise_pred
+        images_rgb = self.decode_with_residual(denoised_latents, *encode_residual)
+        
+        if output_type=='pil':
+            images_rgb = self.image_processor.postprocess(images_rgb)
+        elif output_type=='np':
+            images_rgb = self.image_processor.postprocess(images_rgb, 'np')
+        elif output_type=='pt':
+            images_rgb = self.image_processor.postprocess(images_rgb, 'pt')
+        elif output_type=='none':
+            images_rgb = images_rgb
+        else:
+            raise ValueError('unsupported output_type')
+
+        return images_rgb
+
+    def encode_with_residual(self, sample):
+        re = self.vae.encoder.conv_in(sample)
+        re0, re0_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[0], re)
+        re1, re1_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[1], re0)
+        re2, re2_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[2], re1)
+        re3, re3_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[3], re2)
+        rem = self.vae.encoder.mid_block(re3)
+        re_out = self.vae.encoder.conv_norm_out(rem)
+        re_out = self.vae.encoder.conv_act(re_out)
+        re_out = self.vae.encoder.conv_out(re_out)
+        re_out = self.vae.quant_conv(re_out)
+        
+        posterior = DiagonalGaussianDistribution(re_out)
+        return posterior, (re0_out, re1_out, re2_out, rem, re_out)
+
+    def decode_with_residual(self, z, re0_out, re1_out, re2_out, rem, re_out):
+        rd = self.vae.post_quant_conv(self.se_paths[0](re_out, z))
+        rd = self.vae.decoder.conv_in(rd)
+        rdm = self.vae.decoder.mid_block(self.se_paths[1](rem, rd)).to(torch.float32)
+        rd0 = self.vae.decoder.up_blocks[0](rdm)
+        rd1 = self.vae.decoder.up_blocks[1](self.se_paths[2](re2_out, rd0))
+        rd2 = self.vae.decoder.up_blocks[2](self.se_paths[3](re1_out, rd1))
+        rd3 = self.vae.decoder.up_blocks[3](self.se_paths[4](re0_out, rd2))
+        rd_out = self.vae.decoder.conv_norm_out(rd3)
+        rd_out = self.vae.decoder.conv_act(rd_out)
+        sample_out = self.vae.decoder.conv_out(rd_out)
+        return sample_out
+
+    def _DownEncoderBlock2D_res_forward(self, down_encoder_block_2d, hidden_states):
+        for resnet in down_encoder_block_2d.resnets:
+            hidden_states = resnet(hidden_states, temb=None)
+    
+        output_states = hidden_states
+        if down_encoder_block_2d.downsamplers is not None:
+            for downsampler in down_encoder_block_2d.downsamplers:
+                hidden_states = downsampler(hidden_states)
+    
+        return hidden_states, output_states
+        
+

modeling_sd_gray_inpaint.py

@@ -0,0 +1,98 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from copy import deepcopy
+from torchvision.transforms.functional import rgb_to_grayscale
+import segmentation_models_pytorch as smp
+from diffusers import StableDiffusionInpaintPipeline
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import PretrainedConfig, PreTrainedModel
+
+class SDGrayInpaintConfig(PretrainedConfig):
+    model_type = "sd_gray_inpaint"
+    def __init__(
+        self,
+        base_model="stabilityai/stable-diffusion-2-inpainting",
+        height=512,
+        width=512,
+        **kwargs
+    ):
+        self.base_model=base_model
+        self.height=height
+        self.width=width
+        super().__init__(**kwargs)
+
+class SDGrayInpaintModel(PreTrainedModel):
+    config_class = SDGrayInpaintConfig
+    def __init__(self, config):
+        super().__init__(config)
+        pipe = StableDiffusionInpaintPipeline.from_pretrained(config.base_model)
+        self.mask_predictor = smp.Unet(
+            encoder_name="mit_b4",        
+            encoder_weights="imagenet",     
+            in_channels=3,                  
+            classes=1,                      
+        )
+        self.image_processor = pipe.image_processor
+        self.scheduler = pipe.scheduler
+        self.unet = pipe.unet
+        self.vae = pipe.vae
+        self.prompt_embeds = nn.Parameter(torch.randn(1,77,1024))
+        self.height=config.height
+        self.width=config.width
+
+    def forward(
+        self, 
+        images_gray_masked, 
+        masks=None, 
+        num_inference_steps=250,
+        seed=42,
+        input_type='pil', 
+        output_type='pil'
+    ):
+        generator = torch.Generator()
+        generator.manual_seed(seed)
+        if input_type=='pil':
+            images_gray_masked = self.image_processor.process(images_gray_masked, height=self.height, width=self.width).float()
+        elif input_type=='pt':
+            images_gray_masked=images_gray_masked
+        else:
+            raise ValueError('unsupported input_type')
+        images_gray_masked = images_gray_masked.to(self.vae.device)
+        if masks is None:
+            masks_logits = self.mask_predictor(images_gray_masked)
+            masks = (torch.sigmoid(masks_logits)>0.5)*1.
+        masks = masks.float().to(self.vae.device)
+        B, C, H, W = images_gray_masked.shape
+        prompt_embeds = self.prompt_embeds.repeat(B,1,1)
+
+        scheduler = deepcopy(self.scheduler)
+        scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=self.vae.device)
+        masked_image_latents = self.vae.encode(images_gray_masked).latent_dist.mode() * self.vae.config.scaling_factor
+        mask_latents = F.interpolate(masks, size=(self.unet.config.sample_size, self.unet.config.sample_size))
+        latents = randn_tensor(masked_image_latents.shape, generator=generator).to(self.device) * self.scheduler.init_noise_sigma
+        for t in scheduler.timesteps:
+            latents = scheduler.scale_model_input(latents, t)
+            latent_model_input = torch.cat([latents, mask_latents, masked_image_latents], dim=1)
+            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds)[0]
+            latents = scheduler.step(noise_pred, t, latents)[0]
+        latents = latents / self.vae.config.scaling_factor
+        images_gray_restored = self.vae.decode(latents.detach())[0]
+        images_gray_restored = images_gray_masked * (1-masks) + images_gray_restored.detach() * masks
+        images_gray_restored = rgb_to_grayscale(images_gray_restored)
+        
+        if output_type=='pil':
+            images_gray_restored = self.image_processor.postprocess(images_gray_restored)
+        elif output_type=='np':
+            images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'np')
+        elif output_type=='pt':
+            images_gray_restored = self.image_processor.postprocess(images_gray_restored, 'pt')
+        elif output_type=='none':
+            images_gray_restored = images_gray_restored
+        else:
+            raise ValueError('unsupported output_type')
+
+        return images_gray_restored
+
+

modeling_seresvae.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+
+from diffusers import AutoencoderKL, UNet2DConditionModel
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from transformers import PretrainedConfig, PreTrainedModel
+
+class SEPath(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction=16):
+        super(SEPath, self).__init__()
+        self.fc = nn.Sequential(
+            nn.Linear(in_channels, in_channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_channels // reduction, out_channels, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, in_tensor, out_tensor):
+        B, C, H, W = in_tensor.size()
+        # Squeeze operation
+        x = in_tensor.view(B, C, -1).mean(dim=2)
+        # Excitation operation
+        x = self.fc(x).unsqueeze(2).unsqueeze(2)
+
+        return out_tensor * x
+
+class SeResVaeConfig(PretrainedConfig):
+    model_type = "seresvae"
+    def __init__(
+        self,
+        base_model="stabilityai/stable-diffusion-2-1",
+        height=512,
+        width=512,
+        **kwargs
+    ):
+        self.base_model=base_model
+        self.height=height
+        self.width=width
+        super().__init__(**kwargs)
+
+class SeResVaeModel(PreTrainedModel):
+    config_class = SeResVaeConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.image_processor = VaeImageProcessor()
+        self.vae = AutoencoderKL.from_pretrained(config.base_model, subfolder='vae')
+        self.unet = UNet2DConditionModel.from_pretrained(config.base_model, subfolder='unet')
+        self.se_paths = nn.ModuleList([SEPath(8,4), SEPath(512,512), SEPath(512,512), SEPath(256,512), SEPath(128,256)])
+        self.prompt_embeds = nn.Parameter(torch.randn(1,77,1024))
+        self.height=config.height
+        self.width=config.width
+
+    def forward(self, images_gray, input_type='pil', output_type='pil'):
+        if input_type=='pil':
+            images_gray = self.image_processor.process(images_gray, height=self.height, width=self.width).float()
+        elif input_type=='pt':
+            images_gray=images_gray
+        else:
+            raise ValueError('unsupported input_type')
+        images_gray = images_gray.to(self.vae.device)
+        B, C, H, W = images_gray.shape
+        prompt_embeds = self.prompt_embeds.repeat(B,1,1)
+
+        posterior, encode_residual = self.encode_with_residual(images_gray)
+        latents = posterior.mode()
+        t = torch.LongTensor([500]).repeat(B).to(self.vae.device)
+        noise_pred = self.unet(latents, t, encoder_hidden_states=prompt_embeds)[0]
+        denoised_latents = latents - noise_pred
+        images_rgb = self.decode_with_residual(denoised_latents, *encode_residual)
+        
+        if output_type=='pil':
+            images_rgb = self.image_processor.postprocess(images_rgb)
+        elif output_type=='np':
+            images_rgb = self.image_processor.postprocess(images_rgb, 'np')
+        elif output_type=='pt':
+            images_rgb = self.image_processor.postprocess(images_rgb, 'pt')
+        elif output_type=='none':
+            images_rgb = images_rgb
+        else:
+            raise ValueError('unsupported output_type')
+
+        return images_rgb
+
+    def encode_with_residual(self, sample):
+        re = self.vae.encoder.conv_in(sample)
+        re0, re0_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[0], re)
+        re1, re1_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[1], re0)
+        re2, re2_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[2], re1)
+        re3, re3_out = self._DownEncoderBlock2D_res_forward(self.vae.encoder.down_blocks[3], re2)
+        rem = self.vae.encoder.mid_block(re3)
+        re_out = self.vae.encoder.conv_norm_out(rem)
+        re_out = self.vae.encoder.conv_act(re_out)
+        re_out = self.vae.encoder.conv_out(re_out)
+        re_out = self.vae.quant_conv(re_out)
+        
+        posterior = DiagonalGaussianDistribution(re_out)
+        return posterior, (re0_out, re1_out, re2_out, rem, re_out)
+
+    def decode_with_residual(self, z, re0_out, re1_out, re2_out, rem, re_out):
+        rd = self.vae.post_quant_conv(self.se_paths[0](re_out, z))
+        rd = self.vae.decoder.conv_in(rd)
+        rdm = self.vae.decoder.mid_block(self.se_paths[1](rem, rd)).to(torch.float32)
+        rd0 = self.vae.decoder.up_blocks[0](rdm)
+        rd1 = self.vae.decoder.up_blocks[1](self.se_paths[2](re2_out, rd0))
+        rd2 = self.vae.decoder.up_blocks[2](self.se_paths[3](re1_out, rd1))
+        rd3 = self.vae.decoder.up_blocks[3](self.se_paths[4](re0_out, rd2))
+        rd_out = self.vae.decoder.conv_norm_out(rd3)
+        rd_out = self.vae.decoder.conv_act(rd_out)
+        sample_out = self.vae.decoder.conv_out(rd_out)
+        return sample_out
+
+    def _DownEncoderBlock2D_res_forward(self, down_encoder_block_2d, hidden_states):
+        for resnet in down_encoder_block_2d.resnets:
+            hidden_states = resnet(hidden_states, temb=None)
+    
+        output_states = hidden_states
+        if down_encoder_block_2d.downsamplers is not None:
+            for downsampler in down_encoder_block_2d.downsamplers:
+                hidden_states = downsampler(hidden_states)
+    
+        return hidden_states, output_states
+        
+