Update app.py

app.py

@@ -3,50 +3,51 @@ import torch
 from transformers import CLIPProcessor, CLIPModel
 from torch import nn
 import numpy as np
-import PIL
 from PIL import Image
-from torchvision import transforms
+import torchvision.transforms as transforms
 
 # Load CLIP model and processor
 model_name = "openai/clip-vit-base-patch16"
 clip_model = CLIPModel.from_pretrained(model_name)
 clip_processor = CLIPProcessor.from_pretrained(model_name)
 
-# Generate a random noise tensor (this will be transformed into an image)
+# Define a simple generator network
+class SimpleGenerator(nn.Module):
+    def __init__(self):
+        super(SimpleGenerator, self).__init__()
+        self.fc = nn.Sequential(
+            nn.Linear(512, 1024),
+            nn.ReLU(),
+            nn.Linear(1024, 256*256*3),  # Output image pixels
+            nn.Tanh()  # Normalize output between -1 and 1
+        )
+
+    def forward(self, z):
+        x = self.fc(z)
+        x = x.view(256, 256, 3)  # Reshape to image format
+        return x
+
+# Initialize the generator model
+generator = SimpleGenerator()
+
+# Function to generate an image based on text input
 def generate_image_from_text(text_input):
-    # Preprocess the input text for CLIP model
-    inputs = clip_processor(text=text_input, return_tensors="pt", padding=True)
-
-    # Extract image-text features using CLIP
+    # Preprocess text input using CLIP
+    inputs = clip_processor(text=[text_input], return_tensors="pt", padding=True)
     text_features = clip_model.get_text_features(**inputs)
 
-    # Create a simple GAN-like generator using a random noise tensor
-    class SimpleGenerator(nn.Module):
-        def __init__(self):
-            super(SimpleGenerator, self).__init__()
-            self.fc = nn.Linear(512, 256*256*3)  # Adjust output size to match image dimensions
-            self.relu = nn.ReLU()
-
-        def forward(self, z):
-            x = self.fc(z)
-            x = self.relu(x)
-            x = x.view(-1, 3, 256, 256)  # Reshape to match image shape
-            return x
-
-    # Initialize the generator
-    generator = SimpleGenerator()
-
-    # Generate random noise based on the text features
-    random_input = torch.randn(1, 512)  # Matching CLIP output size (text_features shape)
-    generated_image_tensor = generator(random_input)
-    
-    # Convert generated image tensor to PIL Image
-    generated_image = generated_image_tensor.squeeze().permute(1, 2, 0).detach().numpy()
-    generated_image = np.clip(generated_image, 0, 1)  # Normalize pixel values
-    generated_image = (generated_image * 255).astype(np.uint8)
-    generated_image = Image.fromarray(generated_image)
-
-    return generated_image
+    # Generate image tensor
+    with torch.no_grad():
+        generated_image_tensor = generator(text_features)
+
+    # Normalize tensor to (0, 255)
+    generated_image = (generated_image_tensor - generated_image_tensor.min()) / (generated_image_tensor.max() - generated_image_tensor.min())
+    generated_image = (generated_image * 255).cpu().numpy().astype(np.uint8)
+
+    # Convert to PIL Image
+    image = Image.fromarray(generated_image)
+
+    return image
 
 # Gradio interface
 iface = gr.Interface(fn=generate_image_from_text, inputs="text", outputs="image", live=True)