Update app.py

app.py

@@ -8,18 +8,14 @@ import pandas as pd
 from datasets import load_dataset
 from torch.utils.data import DataLoader, Dataset
 from sklearn.preprocessing import LabelEncoder
-import matplotlib.pyplot as plt
-from PIL import Image
-import io
 
 # Load dataset
-dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:1000]')
+dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
 
 # Preprocess text data
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-
 class CustomDataset(Dataset):
-    def __init__(self, dataset):
+    def init(self, dataset):
         self.dataset = dataset
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
@@ -27,83 +23,65 @@ class CustomDataset(Dataset):
         ])
         self.label_encoder = LabelEncoder()
         self.labels = self.label_encoder.fit_transform(dataset['Model'])
-
-    def __len__(self):
+    def len(self):
         return len(self.dataset)
-
-    def __getitem__(self, idx):
+    def getitem(self, idx):
         image = self.transform(self.dataset[idx]['image'])
         text = tokenizer(self.dataset[idx]['prompt'], padding='max_length', truncation=True, return_tensors='pt')
         label = self.labels[idx]
         return image, text, label
-
+        
 # Define CNN for image processing
 class ImageModel(nn.Module):
-    def __init__(self):
-        super(ImageModel, self).__init__()
+    def init(self):
+        super(ImageModel, self).init()
         self.model = models.resnet18(pretrained=True)
         self.model.fc = nn.Linear(self.model.fc.in_features, 512)
-
     def forward(self, x):
         return self.model(x)
-
+        
 # Define MLP for text processing
 class TextModel(nn.Module):
-    def __init__(self):
-        super(TextModel, self).__init__()
+    def init(self):
+        super(TextModel, self).init()
         self.bert = BertModel.from_pretrained('bert-base-uncased')
         self.fc = nn.Linear(768, 512)
-
     def forward(self, x):
-        output = self.bert(**x)
+        output = self.bert(x)
         return self.fc(output.pooler_output)
-
+        
 # Combined model
 class CombinedModel(nn.Module):
-    def __init__(self):
-        super(CombinedModel, self).__init__()
+    def init**(self):
+        super(CombinedModel, self).init()
         self.image_model = ImageModel()
         self.text_model = TextModel()
         self.fc = nn.Linear(1024, len(dataset['Model']))
-
     def forward(self, image, text):
         image_features = self.image_model(image)
         text_features = self.text_model(text)
         combined = torch.cat((image_features, text_features), dim=1)
         return self.fc(combined)
-
+        
 # Instantiate model
 model = CombinedModel()
-
-def display_recommendations(image, recommended_models, distances):
-    fig, axes = plt.subplots(1, len(recommended_models), figsize=(16, 4))
-    fig.suptitle("Recommended Models")
-
-    for i, (model, distance) in enumerate(zip(recommended_models, distances)):
-        # Load and display the recommended model image
-        model_image = Image.open(io.BytesIO(dataset.get_example(model)['image']))
-        axes[i].imshow(model_image)
-        axes[i].axis('off')
-        axes[i].set_title(f"{model}\nDistance: {distance:.2f}")
-
-    return fig
-
+# Define predict function
 def predict(image):
     model.eval()
     with torch.no_grad():
         image = transforms.ToTensor()(image).unsqueeze(0)
         image = transforms.Resize((224, 224))(image)
         text_input = tokenizer("Sample prompt", return_tensors='pt', padding=True, truncation=True)
-        output = model(image, text_input)
-        distances, indices = torch.topk(-output.squeeze(), 5)
-        recommended_models = [dataset['Model'][i] for i in indices]
-        distances = (-distances).tolist()
-    return display_recommendations(image, recommended_models, distances)
-
+        output = model(image, textinput)
+        , indices = torch.topk(output, 5)
+        recommended_models = [dataset['Model'][i] for i in indices[0]]
+    return recommended_models
+    
+# Set up Gradio interface
 interface = gr.Interface(fn=predict, 
                          inputs=gr.Image(type="pil"), 
-                         outputs=gr.Plot(),
+                         outputs=gr.Textbox(label="Recommended Models"),
                          title="AI Image Model Recommender",
                          description="Upload an AI-generated image to receive model recommendations.")
-
+# Launch the app
 interface.launch()