Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,134 @@
+import gradio as gr
+from transformers import XLNetTokenizer, XLNetModel
+import torch
+import torch.nn as nn
+import numpy as np
+
+# TextEncoder class
+class TextEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = XLNetModel.from_pretrained("xlnet-base-cased")
+
+    def forward(self, input_ids, token_type_ids, attention_mask):
+        hidden = self.transformer(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask).last_hidden_state
+        context = hidden.mean(dim=1)
+        context = context.view(*context.shape, 1, 1)
+        return context
+
+# Generator class
+class Generator(nn.Module):
+    def __init__(self, nz=100, ngf=64, nt=768, nc=3):
+        super().__init__()
+        self.layer1 = nn.Sequential(
+            nn.ConvTranspose2d(nz+nt, ngf*8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(ngf*8)
+        )
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(ngf*8, ngf*2, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf*2),
+            nn.ReLU(True)
+        )
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(ngf*2, ngf*2, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf*2),
+            nn.ReLU(True)
+        )
+        self.layer4 = nn.Sequential(
+            nn.Conv2d(ngf*2, ngf*8, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf*8),
+            nn.ReLU(True)
+        )
+        self.layer5 = nn.Sequential(
+            nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf*4),
+            nn.ReLU(True)
+        )
+        self.layer6 = nn.Sequential(
+            nn.Conv2d(ngf*4, ngf, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True)
+        )
+        self.layer7 = nn.Sequential(
+            nn.Conv2d(ngf, ngf, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True)
+        )
+        self.layer8 = nn.Sequential(
+            nn.Conv2d(ngf, ngf*4, 3, 1, 1),
+            nn.Dropout2d(inplace=True),
+            nn.BatchNorm2d(ngf*4),
+            nn.ReLU(True)
+        )
+        self.layer9 = nn.Sequential(
+            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf*2),
+            nn.ReLU(True)
+        )
+        self.layer10 = nn.Sequential(
+            nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True)
+        )
+        self.layer11 = nn.Sequential(
+            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
+            nn.Tanh()
+        )
+
+    def forward(self, noise, encoded_text):
+        x = torch.cat([noise, encoded_text], dim=1)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.layer5(x)
+        x = self.layer6(x)
+        x = self.layer7(x)
+        x = self.layer8(x)
+        x = self.layer9(x)
+        x = self.layer10(x)
+        x = self.layer11(x)
+        return x
+
+
+# Load the model and tokenizer
+model_path = "./checkpoint.pth"
+tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
+text_encoder = XLNetModel.from_pretrained('xlnet-base-cased')
+model = Generator()
+model_state_dict = torch.load(model_path, map_location="cpu")
+generator = model_state_dict['models']['generator']
+model.load_state_dict(generator)
+
+text_encoder.to("cpu")
+model.to("cpu")
+model.eval()
+
+# Functions to encode text and generate image
+def encode_text(text):
+    text_encoder_model = TextEncoder()
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    encoded_text = text_encoder_model(**inputs)
+    return encoded_text
+
+def generate_image(text):
+    encoded_text = encode_text(text)
+    noise = torch.randn((1, 100, 1, 1), device="cpu")
+    with torch.no_grad():
+        generated_image = model(noise, encoded_text).detach().squeeze().cpu()
+    gen_image_np = generated_image.numpy()
+    gen_image_np = np.transpose(gen_image_np, (1, 2, 0))  # Change from CHW to HWC
+    gen_image_np = (gen_image_np - gen_image_np.min()) / (gen_image_np.max() - gen_image_np.min())  # Normalize to [0, 1]
+    gen_image_np = (gen_image_np * 255).astype(np.uint8)
+    return gen_image_np
+
+# Gradio interface
+inputs = gr.inputs.Textbox(label="Enter a flower-related description", default="A beautiful red rose")
+outputs = gr.outputs.Image(type="numpy", label="Generated Flower Image")
+
+gr.Interface(fn=generate_image, inputs=inputs, outputs=outputs, title="Flower Image Generator", description="Enter a description of a flower to generate an image.").launch()