Uploading logic files

README.md

@@ -1,14 +1,13 @@
 ---
-title: Attndecgan
-emoji: 😻
+title: Attndecgan Demo
+emoji: 🔥
 colorFrom: gray
-colorTo: indigo
+colorTo: green
 sdk: gradio
-sdk_version: 6.5.0
+sdk_version: 5.41.1
 app_file: app.py
 pinned: false
-license: mit
-short_description: Att-DeCGAN demo
+short_description: A demo for the attention diversity enhanced GAN
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,135 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+from huggingface_hub import hf_hub_download
+import os
+
+# Import your networks (you'll need to upload networks.py to your Space)
+from networks import HPBGenerator  # Adjust this import based on your networks.py structure
+
+class CycleGANInference:
+    def __init__(self, model_repo_id, checkpoint_filename_AtoB, checkpoint_filename_BtoA=None):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Download model checkpoints from Hugging Face Hub
+        checkpoint_path_AtoB = hf_hub_download(
+            repo_id=model_repo_id, 
+            filename=checkpoint_filename_AtoB
+        )
+        
+        # Initialize generators
+        # Adjust these parameters based on your model architecture
+        self.netG_A2B = HPBGenerator(input_nc=3, output_nc=3, ngf=64, n_blocks=9)  # A to B
+        
+        if checkpoint_filename_BtoA:
+            checkpoint_path_BtoA = hf_hub_download(
+                repo_id=model_repo_id, 
+                filename=checkpoint_filename_BtoA
+            )
+            self.netG_B2A = HPBGenerator(input_nc=3, output_nc=3, ngf=64, n_blocks=9)  # B to A
+        else:
+            self.netG_B2A = None
+        
+        # Load model weights
+        self.netG_A2B.load_state_dict(torch.load(checkpoint_path_AtoB, map_location=self.device))
+        if self.netG_B2A and checkpoint_filename_BtoA:
+            self.netG_B2A.load_state_dict(torch.load(checkpoint_path_BtoA, map_location=self.device))
+        
+        # Set to evaluation mode
+        self.netG_A2B.eval()
+        if self.netG_B2A:
+            self.netG_B2A.eval()
+        
+        # Move to device
+        self.netG_A2B.to(self.device)
+        if self.netG_B2A:
+            self.netG_B2A.to(self.device)
+        
+        # Define transforms
+        self.transform = transforms.Compose([
+            transforms.Resize((256, 256)),  # Adjust size based on your model
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+        
+        self.inverse_transform = transforms.Compose([
+            transforms.Normalize((-1, -1, -1), (2, 2, 2)),  # Denormalize
+            transforms.ToPILImage()
+        ])
+    
+    def transform_image(self, image, direction="A_to_B"):
+        # Preprocess
+        input_tensor = self.transform(image).unsqueeze(0).to(self.device)
+        
+        with torch.no_grad():
+            if direction == "A_to_B":
+                output_tensor = self.netG_A2B(input_tensor)
+            elif direction == "B_to_A" and self.netG_B2A:
+                output_tensor = self.netG_B2A(input_tensor)
+            else:
+                raise ValueError("Invalid direction or model not available")
+        
+        # Postprocess
+        output_image = self.inverse_transform(output_tensor.squeeze(0).cpu())
+        return output_image
+
+# Initialize your model
+# Replace these with your actual Hugging Face repo ID and checkpoint filenames
+MODEL_REPO_ID = "profmatthew/Attn-DeCGAN"  # Replace with your repo
+CHECKPOINT_A2B = "latest_net_G_A.pth"  # Replace with your checkpoint filename
+CHECKPOINT_B2A = "latest_net_G_B.pth"  # Replace with your checkpoint filename (optional)
+
+cyclegan_model = CycleGANInference(
+    model_repo_id=MODEL_REPO_ID,
+    checkpoint_filename_AtoB=CHECKPOINT_A2B,
+    checkpoint_filename_BtoA=CHECKPOINT_B2A  # Set to None if you only have one direction
+)
+
+def generate_image(input_image, direction):
+    try:
+        output_image = cyclegan_model.transform_image(input_image, direction)
+        return output_image
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="CycleGAN Image Translation") as demo:
+    gr.Markdown("# CycleGAN Image Translation")
+    gr.Markdown("Upload an image and select the transformation direction.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            direction = gr.Dropdown(
+                choices=["A_to_B", "B_to_A"], 
+                value="A_to_B", 
+                label="Translation Direction"
+            )
+            generate_btn = gr.Button("Generate", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Generated Image")
+    
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[input_image, direction],
+        outputs=output_image
+    )
+    
+    # Add some examples if you have them
+    # gr.Examples(
+    #     examples=[
+    #         # Add paths to example images here
+    #         # ["example1.jpg", "A_to_B"],
+    #         # ["example2.jpg", "B_to_A"],
+    #     ],
+    #     inputs=[input_image, direction],
+    #     outputs=output_image,
+    #     fn=generate_image,
+    # )
+
+if __name__ == "__main__":
+    demo.launch()

gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

networks.py

@@ -0,0 +1,394 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+import torch.nn.functional as F
+from torch import nn, einsum
+from einops import rearrange, reduce, repeat
+
+###############################################################################
+# Helper Functions
+###############################################################################
+
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+    elif norm_type == 'none':
+        norm_layer = None
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def get_scheduler(optimizer, opt):
+    if opt.lr_policy == 'lambda':
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.lr_policy == 'step':
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
+    elif opt.lr_policy == 'plateau':
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif opt.lr_policy == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0)
+    else:
+        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
+    return scheduler
+
+
+def init_weights(net, init_type='normal', gain=0.02):
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        net.to(gpu_ids[0])
+        net = torch.nn.DataParallel(net, gpu_ids)
+    init_weights(net, init_type, gain=init_gain)
+    return net
+
+
+
+
+##############################################################################
+# Classes
+##############################################################################
+
+
+# Defines the GAN loss which uses either LSGAN or the regular GAN.
+# When LSGAN is used, it is basically same as MSELoss,
+# but it abstracts away the need to create the target label tensor
+# that has the same size as the input
+class GANLoss(nn.Module):
+    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0):
+        super(GANLoss, self).__init__()
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        if use_lsgan:
+            self.loss = nn.MSELoss()
+        else:
+            self.loss = nn.BCELoss()
+    def get_target_tensor(self, input, target_is_real):
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(input)
+
+    def __call__(self, input, target_is_real):
+        target_tensor = self.get_target_tensor(input, target_is_real)
+        return self.loss(input, target_tensor)
+
+#################################################################################
+#                    Critic Loss for Wassertein Gan GP                          #
+#################################################################################
+class GradPenalty(nn.Module):
+    def __init__(self, use_cuda):
+        super(GradPenalty, self).__init__()
+        self.use_cuda = use_cuda
+    def forward(self, critic, real_data, fake_data):
+        alpha = torch.rand_like(real_data)
+
+        assignGPU = lambda x: x.cuda() if self.use_cuda else x
+        alpha = assignGPU(alpha)
+
+        interpolates = alpha*real_data + (1-alpha)*fake_data.detach()
+        interpolates = assignGPU(interpolates)
+        interpolates = torch.autograd.Variable(interpolates, requires_grad = True)
+
+        critic_interpolates = critic(interpolates)
+
+        gradients = torch.autograd.grad(
+            outputs=critic_interpolates, 
+            inputs=interpolates,
+            grad_outputs=assignGPU(torch.ones(critic_interpolates.size())),
+            create_graph=True, retain_graph=True, only_inputs=True
+        )[0]
+        gradients = gradients.view(gradients.size(0), -1)
+        gradient_penalty = ((gradients.norm(2, dim=1)-1)**2).mean()
+        return gradient_penalty
+
+#####
+#####
+
+#################################################################################
+#                   Hybrid Perception Block and DPSA LAyer                      #
+#################################################################################
+
+
+# helper functions
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def l2norm(t):
+    return F.normalize(t, dim = -1)
+
+# helper classes
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + x
+
+class ChanLayerNorm(nn.Module):
+    def __init__(self, dim, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
+
+    def forward(self, x):
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
+
+# classes
+
+class HPB(nn.Module):
+    """ Hybrid Perception Block """
+
+    def __init__(
+        self,
+        dim,
+        dim_head = 32,
+        heads = 8,
+        ff_mult = 4,
+        attn_height_top_k = 8,
+        attn_width_top_k = 8,
+        attn_dropout = 0.,
+        ff_dropout = 0.
+    ):
+        super().__init__()
+
+        self.attn = DPSA(
+            dim = dim,
+            heads = heads,
+            dim_head = dim_head,
+            height_top_k = attn_height_top_k,
+            width_top_k = attn_width_top_k,
+            dropout = attn_dropout
+        )
+
+        self.dwconv = nn.Conv2d(dim, dim, 3, padding = 1, groups = dim)
+        self.attn_parallel_combine_out = nn.Conv2d(dim * 2, dim, 1)
+
+        ff_inner_dim = dim * ff_mult
+
+        self.ff = nn.Sequential(
+            nn.Conv2d(dim, ff_inner_dim, 1),
+            nn.InstanceNorm2d(ff_inner_dim),
+            nn.GELU(),
+            nn.Dropout(ff_dropout),
+            Residual(nn.Sequential(
+                nn.Conv2d(ff_inner_dim, ff_inner_dim, 3, padding = 1, groups = ff_inner_dim),
+                nn.InstanceNorm2d(ff_inner_dim),
+                nn.GELU(),
+                nn.Dropout(ff_dropout)
+            )),
+            nn.Conv2d(ff_inner_dim, dim, 1),
+            nn.InstanceNorm2d(ff_inner_dim)
+        )
+
+    def forward(self, x):
+        attn_branch_out = self.attn(x)
+        conv_branch_out = self.dwconv(x)
+
+        concatted_branches = torch.cat((attn_branch_out, conv_branch_out), dim = 1)
+        attn_out = self.attn_parallel_combine_out(concatted_branches) + x
+
+        return self.ff(attn_out)
+
+class DPSA(nn.Module):
+    """ Dual-pruned Self-attention Block """
+
+    def __init__(
+        self,
+        dim,
+        height_top_k = 8,
+        width_top_k = 8,
+        dim_head = 32,
+        heads = 8,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.heads = heads
+        self.dim_head = dim_head
+        self.scale = dim_head ** -0.5
+        inner_dim = heads * dim_head
+
+        self.norm = ChanLayerNorm(dim)
+        self.to_qkv = nn.Conv2d(dim, inner_dim * 3, 1, bias = False)
+
+        self.height_top_k = height_top_k
+        self.width_top_k = width_top_k
+
+        self.dropout = nn.Dropout(dropout)
+        self.to_out = nn.Conv2d(inner_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        x = self.norm(x)
+
+        q, k, v = self.to_qkv(x).chunk(3, dim = 1)
+
+        # fold out heads
+
+        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> (b h) c x y', h = self.heads), (q, k, v))
+
+        # they used l2 normalized queries and keys, cosine sim attention basically
+
+        q, k = map(l2norm, (q, k))
+
+        # calculate whether to select and rank along height and width
+
+        need_height_select_and_rank = self.height_top_k < h
+        need_width_select_and_rank = self.width_top_k < w
+
+        # select and rank keys / values, probing with query (reduced along height and width) and keys reduced along row and column respectively
+
+        if need_width_select_and_rank or need_height_select_and_rank:
+            q_probe = reduce(q, 'b h w d -> b d', 'sum')
+
+        # gather along height, then width
+
+        if need_height_select_and_rank:
+            k_height = reduce(k, 'b h w d -> b h d', 'sum')
+
+            top_h_indices = einsum('b d, b h d -> b h', q_probe, k_height).topk(k = self.height_top_k, dim = -1).indices
+
+            top_h_indices = repeat(top_h_indices, 'b h -> b h w d', d = self.dim_head, w = k.shape[-2])
+
+            k, v = map(lambda t: t.gather(1, top_h_indices), (k, v)) # first gather across height
+
+        if need_width_select_and_rank:
+            k_width = reduce(k, 'b h w d -> b w d', 'sum')
+
+            top_w_indices = einsum('b d, b w d -> b w', q_probe, k_width).topk(k = self.width_top_k, dim = -1).indices
+
+            top_w_indices = repeat(top_w_indices, 'b w -> b h w d', d = self.dim_head, h = k.shape[1])
+
+            k, v = map(lambda t: t.gather(2, top_w_indices), (k, v)) # then gather along width
+
+        # select the appropriate keys and values
+
+        q, k, v = map(lambda t: rearrange(t, 'b ... d -> b (...) d'), (q, k, v))
+
+        # cosine similarities
+
+        sim = einsum('b i d, b j d -> b i j', q, k)
+
+        # attention
+
+        attn = sim.softmax(dim = -1)
+        attn = self.dropout(attn)
+
+        # aggregate out
+
+        out = einsum('b i j, b j d -> b i d', attn, v)
+
+        # merge heads and combine out
+
+        out = rearrange(out, '(b h) (x y) d -> b (h d) x y', x = h, y = w, h = self.heads)
+        return self.to_out(out)
+
+#####
+#####
+
+# New HybridPerceptionBlockGenerator
+
+class HPBGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.InstanceNorm2d, use_dropout=False, n_blocks=9, padding_type='reflect'):
+        assert(n_blocks >= 0)
+        super(HPBGenerator, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [
+            nn.ReflectionPad2d(3),
+                nn.Conv2d(
+                input_nc, ngf, 
+                kernel_size=7, 
+                padding=0,
+                bias=use_bias
+            ),
+            norm_layer(ngf),
+            nn.GELU()
+        ]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [
+                nn.Conv2d(
+                    ngf * mult, ngf * mult * 2, kernel_size=3,
+                    stride=2, padding=1, bias=use_bias
+                ),
+                norm_layer(ngf * mult * 2),
+                nn.GELU()
+            ]
+
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [
+                HPB(ngf * mult, ngf)
+            ]
+
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [
+                nn.ConvTranspose2d(
+                    ngf * mult, int(ngf * mult / 2),
+                    kernel_size=3, stride=2,
+                    padding=1, output_padding=1,
+                    bias=use_bias
+                ),
+                norm_layer(int(ngf * mult / 2)),
+                nn.GELU()
+            ]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+torchvision
+Pillow
+numpy
+huggingface_hub
+einops