Upload app.py

app.py

@@ -1,30 +1,15 @@
-max_description_length=100
-
-
-
-
-
-
 from skimage import color
 import numpy as np
 
-def hex_to_rgb(hex_color):
-    hex_color = hex_color.lstrip('#')
-    return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))
+tokenizer_input_length = 77
+
+import torch
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def rgb_to_hex(rgb_array):
-    return "#{:02x}{:02x}{:02x}".format(*rgb_array)
+    return "{:02x}{:02x}{:02x}".format(*rgb_array)
 
-def rgb_to_normalized_lab(rgb_array):
-    rgb_array = np.array(rgb_array, dtype=np.float32) / 255.0
-    
-    if rgb_array.ndim == 1:
-        rgb_array = rgb_array.reshape(1, 3)
 
-    lab_array = color.rgb2lab(rgb_array)
-    lab_array[:,0] /= 100.0
-    lab_array[:,1:] /= 127.0 
-    return tuple(lab_array.squeeze())
 
 def normalized_lab_to_rgb(lab_array):
     lab_array = np.array(lab_array, dtype=np.float32)
@@ -42,19 +27,19 @@ def normalized_lab_to_rgb(lab_array):
     return tuple(rgb_array.squeeze())
     
     
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers import RobertaModel, RobertaTokenizer
+
 from huggingface_hub import hf_hub_download
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-print(f'Using {device}')
+model_path = hf_hub_download(repo_id="lasercatz/text2palette", filename="epoch_19.pth")
+
 
-model_path = hf_hub_download(repo_id="lasercatz/text2palette", filename="epoch_13.pth")
 
 
 
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import CLIPTextModel, CLIPTokenizer
+
 
 class AttentionPooling(nn.Module):
     def __init__(self, d_model):
@@ -68,6 +53,7 @@ class AttentionPooling(nn.Module):
         weights = F.softmax(scores, dim=-1).unsqueeze(-1)
         return torch.sum(x * weights, dim=1)
 
+
 class SequencePriorNet(nn.Module):
     def __init__(self, d_model, d_z, n_heads=4):
         super().__init__()
@@ -78,53 +64,56 @@ class SequencePriorNet(nn.Module):
         self.dropout = nn.Dropout(0.3)
 
     def forward(self, text_feats, attention_mask):
-        attn_output, _ = self.attn(text_feats, text_feats, text_feats, key_padding_mask=~attention_mask.bool())
+        attn_output, _ = self.attn(
+            text_feats, text_feats, text_feats, key_padding_mask=~attention_mask.bool())
         x = self.norm(attn_output + text_feats)
         x = self.dropout(x)
         x = self.pool(x, attention_mask)
         x = self.fc(x)
         return x
 
-class TextToPaletteModel(nn.Module):
-    def __init__(self, d_model=768, d_z=256, max_text_len=max_description_length, max_seq_len=10, 
+
+class Text2PaletteModel(nn.Module):
+    def __init__(self, d_model=768, d_z=256, max_seq_len=64,
                  n_layers=8, n_heads=8, dim_ff=3072):
         super().__init__()
         self.d_model = d_model
-        self.max_text_len= max_text_len
         self.max_seq_len = max_seq_len
 
-        # Text encoder
-        self.tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
-        self.roberta = RobertaModel.from_pretrained('roberta-base')
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            'openai/clip-vit-base-patch32')
+        self.clip_text = CLIPTextModel.from_pretrained(
+            'openai/clip-vit-base-patch32')
+        
+        self.tokenizer_input_length = tokenizer_input_length
+
+
         self.text_proj = nn.Sequential(
-            nn.Linear(768, d_model*2),
+            nn.Linear(512, d_model*2),
             nn.GELU(),
             nn.LayerNorm(d_model*2),
             nn.Dropout(0.3),
             nn.Linear(d_model*2, d_model)
         )
 
-        # Color processing
         self.color_embed = nn.Sequential(
             nn.Linear(3, d_model),
             nn.LayerNorm(d_model),
             nn.GELU(),
             nn.Dropout(0.3)
         )
-        
+
         self.cross_attn = nn.MultiheadAttention(d_model, 8, batch_first=True)
-        
-        # Positional embeddings
+
         self.position_embed = nn.Embedding(max_seq_len, d_model)
         self.start_embed = nn.Parameter(torch.randn(1, d_model))
 
-        # Palette encoder
         self.palette_encoder = nn.TransformerEncoder(
-            nn.TransformerEncoderLayer(d_model, n_heads, dim_ff, batch_first=True),
+            nn.TransformerEncoderLayer(
+                d_model, n_heads, dim_ff, batch_first=True),
             n_layers
         )
 
-        # Latent projection
         self.z_proj = nn.Sequential(
             nn.Linear(d_model*2, d_z),
             nn.LayerNorm(d_z),
@@ -134,13 +123,12 @@ class TextToPaletteModel(nn.Module):
         self.z_mu = nn.Linear(d_z, d_z)
         self.z_logvar = nn.Linear(d_z, d_z)
 
-        # Decoder
         self.decoder = nn.TransformerDecoder(
-            nn.TransformerDecoderLayer(d_model, n_heads, dim_ff, batch_first=True),
+            nn.TransformerDecoderLayer(
+                d_model, n_heads, dim_ff, batch_first=True),
             n_layers
         )
 
-        # Output layers
         self.out_mu_L = nn.Sequential(
             nn.Linear(d_model, 1),
             nn.Sigmoid()
@@ -150,79 +138,82 @@ class TextToPaletteModel(nn.Module):
             nn.Tanh()
         )
         self.out_logvar = nn.Linear(d_model, 3)
-        
+
         self.prior_net = SequencePriorNet(d_model, d_z, n_heads=4)
 
-        # Pooling
         self.text_pool = AttentionPooling(d_model)
         self.palette_pool = AttentionPooling(d_model)
 
+    def reparameterize(self, mu, logvar):
+
+        if self.training:
+            std = torch.exp(0.5 * logvar)
+            eps = torch.randn_like(std)
+            return mu + eps * std
+        else:
+            return mu
+
 
     @torch.no_grad()
-    def generate(self, text, palette_size,temp=1.0):
+    def generate(self, text, palette_size, temp=1.0):
         self.eval()
-        device = next(self.parameters()).device
-        tokenized = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, 
-                                  max_length=self.max_text_len).to(device)
-        
-        # Text features
-        text_feats = self.roberta(**tokenized).last_hidden_state
+        tokenized = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True,
+                                   max_length=self.tokenizer_input_length).to(next(self.parameters()).device)
+
+        text_feats = self.clip_text(**tokenized).last_hidden_state
         text_feats = self.text_proj(text_feats)
-        text_pooled = self.text_pool(text_feats, tokenized['attention_mask'])
-        
+
         # Sample from prior
         prior_params = self.prior_net(text_feats, tokenized['attention_mask'])
         prior_mu, prior_logvar = prior_params.chunk(2, -1)
-        z = prior_mu + torch.exp(0.5 * prior_logvar) * torch.randn_like(prior_mu) * temp 
+        z = prior_mu + torch.exp(0.5 * prior_logvar) * \
+            torch.randn_like(prior_mu) * temp
         z_expanded = self.z_expand(z).unsqueeze(1)
 
-        memory = torch.cat([z_expanded, text_feats], dim=1)  # [1, 1 + seq_len, d_model]
+        memory = torch.cat([z_expanded, text_feats],
+                           dim=1)  # [1, T+1, d_model]
         memory_key_padding_mask = torch.cat([
             torch.zeros((1, 1), dtype=torch.bool, device=device),
             ~tokenized['attention_mask'].bool()
-        ], dim=1)  # [1, 1 + seq_len]
+        ], dim=1)  # [1, T+1]
 
-        generate_size = min(palette_size, self.max_seq_len)
-    
-
-        # Generation loop
         colors = []
-        current_emb = self.start_embed.unsqueeze(0)
-        
-        for i in range(generate_size):
-            # Positional update
-            pos = self.position_embed(torch.arange(0, current_emb.size(1), device=device)).unsqueeze(0)  # [1, i+1, d_model]
+        batch_size = 1
+        current_emb = self.start_embed.unsqueeze(0).expand(
+            batch_size, -1, -1)  # [1, 1, d_model]
+
+        for i in range(min(palette_size, self.max_seq_len)):
+
+            pos = self.position_embed(torch.arange(0, current_emb.size(
+                1), device=device)).unsqueeze(0)  # [1, i+1, d_model]
             decoder_in = current_emb + pos  # [1, i+1, d_model]
-            
-            # Decode
+
             output = self.decoder(
                 decoder_in,
                 memory,
-                tgt_mask=nn.Transformer.generate_square_subsequent_mask(decoder_in.size(1), device=device),
+                tgt_mask=nn.Transformer.generate_square_subsequent_mask(
+                    decoder_in.size(1), device=device),
                 memory_key_padding_mask=memory_key_padding_mask
             )  # [1, i+1, d_model]
-            
-            # Predict color
-            mu = torch.cat([self.out_mu_L(output[:, -1]), self.out_mu_ab(output[:, -1])], dim=-1)  # [1, 3]
-            logvar = self.out_logvar(output[:, -1])  # [1, 3]
-            
 
-            color = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu) * temp  # [1, 3]
-                
+            mu = torch.cat([self.out_mu_L(output[:, -1]),
+                           self.out_mu_ab(output[:, -1])], dim=-1)  # [1, 3]
+            logvar = self.out_logvar(output[:, -1])  # [1, 3]
+            color = mu + torch.exp(0.5 * logvar) * \
+                torch.randn_like(mu) * temp  # [1, 3]
             color[:, 0].clamp_(0, 1)
             color[:, 1:].clamp_(-1, 1)
-            
+
             colors.append(color)
-            
-            # Embed the color and update the sequence
+
             color_emb = self.color_embed(color.unsqueeze(1))  # [1, 1, d_model]
+            current_emb = torch.cat(
+                [current_emb, color_emb], dim=1)  # [1, i+2, d_model]
 
-            current_emb = torch.cat([current_emb, color_emb ], dim=1)
-        
         return torch.cat(colors, dim=0).unsqueeze(0)
 
 
-model = TextToPaletteModel().to(device)
+model = Text2PaletteModel().to(device)
 state_dict = torch.load(model_path, map_location=torch.device(device))
 model.load_state_dict(state_dict['model'])
 model.to(device)
@@ -231,7 +222,78 @@ model.eval()
 import gradio as gr
 
 
-def generate_palettes(text, palette_size=5, temp=1.0):
+
+def generate_palette(text, palette_size=5, temp=1.0):
+
+    html=""
+
+    with torch.no_grad():
+        generated_palette = model.generate(
+            text, 
+            palette_size=int(palette_size),
+            temp=temp
+        )
+    
+    lab = generated_palette[0].cpu().numpy()
+    hex_palette = [rgb_to_hex(normalized_lab_to_rgb(lab_color)) for lab_color in lab]
+    
+    html += "<div style='display: flex; flex-direction: row;align-items: center; width:100%;'>"
+    
+    hex_codes = []
+    for i,hex_color in enumerate(hex_palette):
+        hex_color = "#"+hex_color.upper()
+        hex_codes.append(hex_color)
+        html += f'<div style=\'margin:0;flex: 1; text-align: center;\'><div style=\'background-color: {hex_color}; width: 100%; height: 100px;border-radius:{"1em 0 0 1em" if i==0 else "0 1em 1em 0" if i==len(hex_palette)-1 else "0"}\'></div><p style=\'font-size: 14px; margin-top: 5px;\'>{hex_color}</p></div>'
+    html += "</div>"
+
+    return html
+
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("<h1>Palette Generator</h1>")
+    
+    input = gr.Textbox(label="Input text", placeholder="Describe the palette in your mind")
+
+    with gr.Row():
+        palette_size = gr.Slider(2, 10, value=5, step=1, label="Colors")
+        temp = gr.Slider(0.0, 0.1, value=0.03, step=0.01, label="Temperature")
+    with gr.Row():
+        with gr.Column():
+            gr.Examples(
+                examples=[["fries in ketchup"], ["blueberry milkshake"], ["Oreo McFlurry"]],
+                inputs=[input],
+                label="Food & Drinks"
+            )
+        with gr.Column():
+            gr.Examples(
+                examples=[["bonfire"], ["sheep on grass"], ["North Arctic"]],
+                inputs=[input],
+                label="Objects & Places"
+            )
+    with gr.Row():
+        with gr.Column():
+            gr.Examples(
+                examples=[["rock climbing"], ["scuba-diving"], ["Halloween pumpkin party"]],
+                inputs=[input],
+                label="Activities"
+            )
+        with gr.Column():
+            gr.Examples(
+                examples=[["sweetheart"], ["sorrow"], ["murder"]],
+                inputs=[input],
+                label="Abstract"
+            )
+    generate_button = gr.Button("🎨 Generate")
+    output = gr.HTML("<div style=\"height: 100px\"></div>")
+
+    generate_button.click(
+        generate_palette,
+        inputs=[input, palette_size, temp],
+        outputs=output
+    )
+
+def generate_palettes_api(text, palette_size=5, temp=1.0):
     temps=[temp]*4
     hex_palettes = []
     
@@ -249,21 +311,4 @@ def generate_palettes(text, palette_size=5, temp=1.0):
         hex_palettes.append(hex_palette)
     return hex_palettes
 
-
-
-
-with gr.Blocks() as demo:
-  default_input = gr.Textbox(label="Input text", placeholder="")
-  palette_size = gr.Slider(1, 10, value=5, step=1, label="Palette size")
-  temp = gr.Slider(0.0, 0.1, value=0.03, step=0.01, label="Temperature")
-  default_button = gr.Button("🎨 Generate")
-  default_output = gr.HTML()
-
-  default_button.click(
-      generate_palettes,
-      inputs=[default_input,palette_size , temp],
-      outputs=default_output
-)
-    
-
 demo.launch()