Update app.py

app.py

@@ -1,80 +1,35 @@
 import gradio as gr
-from transformers import XLNetTokenizer, XLNetModel
 import torch
 import torch.nn as nn
+from transformers import XLNetTokenizer, XLNetModel
 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)
+            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.Conv2d(ngf * 8, ngf * 2, 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.BatchNorm2d(ngf * 2),
             nn.ReLU(True)
         )
+        # Add other layers as needed...
+
         self.layer11 = nn.Sequential(
             nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
             nn.Tanh()
@@ -84,23 +39,16 @@ class Generator(nn.Module):
         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)
+        # Pass through other layers...
         x = self.layer11(x)
         return x
 
-
 # Load the model and tokenizer
-model_path = "./checkpoint.pth"
+model_path = "checkpoint.pth"  # Adjust as necessary
 tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
-text_encoder = XLNetModel.from_pretrained('xlnet-base-cased')
+text_encoder = TextEncoder()
 model = Generator()
+
 model_state_dict = torch.load(model_path, map_location="cpu")
 generator = model_state_dict['models']['generator']
 model.load_state_dict(generator)
@@ -109,26 +57,37 @@ 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)
+def generate_image(enc_text):
     noise = torch.randn((1, 100, 1, 1), device="cpu")
     with torch.no_grad():
-        generated_image = model(noise, encoded_text).detach().squeeze().cpu()
+        generated_image = model(noise, enc_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")
+def encode_text(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    encoded_text = text_encoder(**inputs)
+    return encoded_text
+
+def on_generate_button_click(text_input):
+    if text_input:
+        encoded_text = encode_text(text_input)
+        generated_image = generate_image(encoded_text)
+        return generated_image
+    return None
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Flower Image Generator")
+    text_input = gr.Textbox(label="Enter a flower-related description", value="A beautiful red rose")
+    generate_button = gr.Button("Generate Image")
+    output_image = gr.Image(type="numpy")  # Ensure output type is correct
+
+    generate_button.click(on_generate_button_click, inputs=text_input, outputs=output_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()
+# Launch the Gradio app
+demo.launch()