Update app.py

app.py

@@ -6,39 +6,14 @@ from torchvision import models
 from transformers import BertTokenizer, BertModel
 import pandas as pd
 from datasets import load_dataset
-from torch.utils.data import DataLoader, Dataset, random_split
+from torch.utils.data import DataLoader, Dataset
 from sklearn.preprocessing import LabelEncoder
-from sklearn.metrics import confusion_matrix, classification_report, accuracy_score
-import seaborn as sns
-import matplotlib.pyplot as plt
-import numpy as np
-from tqdm import tqdm
-import os
-import logging
-
-# Set up logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.FileHandler('model_training.log'),
-        logging.StreamHandler()
-    ]
-)
-
-# Create output directory for results
-os.makedirs('output', exist_ok=True)
 
 # Load dataset and filter out null/none values
-logging.info("Loading and filtering dataset...")
 dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
+# Filter out entries where Model is None or empty
 dataset = dataset.filter(lambda example: example['Model'] is not None and example['Model'].strip() != '')
 
-if len(dataset) == 0:
-    raise ValueError("Dataset is empty after filtering!")
-
-logging.info(f"Dataset size after filtering: {len(dataset)}")
-
 # Preprocess text data
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
@@ -48,329 +23,63 @@ class CustomDataset(Dataset):
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
-                               std=[0.229, 0.224, 0.225])
         ])
         self.label_encoder = LabelEncoder()
         self.labels = self.label_encoder.fit_transform(dataset['Model'])
-        self.unique_models = self.label_encoder.classes_
-        
-        logging.info(f"Number of unique models: {len(self.unique_models)}")
 
     def __len__(self):
         return len(self.dataset)
 
     def __getitem__(self, idx):
-        try:
-            image = self.transform(self.dataset[idx]['image'])
-            text = tokenizer(
-                self.dataset[idx]['prompt'],
-                padding='max_length',
-                truncation=True,
-                max_length=512,
-                return_tensors='pt'
-            )
-            label = self.labels[idx]
-            return image, text, label
-        except Exception as e:
-            logging.error(f"Error processing item {idx}: {str(e)}")
-            raise
+        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__()
         self.model = models.resnet18(pretrained=True)
         self.model.fc = nn.Linear(self.model.fc.in_features, 512)
-        
+
     def forward(self, x):
-        x = self.model(x)
-        return nn.functional.relu(x)
+        return self.model(x)
 
+# Define MLP for text processing
 class TextModel(nn.Module):
     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):
-        outputs = self.bert(**x)
-        x = outputs.pooler_output
-        x = self.fc(x)
-        return nn.functional.relu(x)
+        output = self.bert(**x)
+        return self.fc(output.pooler_output)
 
+# Combined model
 class CombinedModel(nn.Module):
-    def __init__(self, num_classes):
+    def __init__(self):
         super(CombinedModel, self).__init__()
         self.image_model = ImageModel()
         self.text_model = TextModel()
-        self.dropout = nn.Dropout(0.2)
-        self.fc = nn.Linear(1024, num_classes)
-        
+        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)
-        combined = self.dropout(combined)
         return self.fc(combined)
 
-class ModelTrainerEvaluator:
-    def __init__(self, model, dataset, batch_size=32, learning_rate=0.001):
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        logging.info(f"Using device: {self.device}")
-        
-        self.model = model.to(self.device)
-        self.batch_size = batch_size
-        self.criterion = nn.CrossEntropyLoss()
-        self.optimizer = torch.optim.AdamW(
-            model.parameters(),
-            lr=learning_rate,
-            weight_decay=0.01
-        )
-        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            self.optimizer,
-            mode='min',
-            factor=0.1,
-            patience=2,
-            verbose=True
-        )
-        
-        # Split dataset
-        total_size = len(dataset)
-        train_size = int(0.7 * total_size)
-        val_size = int(0.15 * total_size)
-        test_size = total_size - train_size - val_size
-        
-        train_dataset, val_dataset, test_dataset = random_split(
-            dataset, [train_size, val_size, test_size]
-        )
-        
-        self.train_loader = DataLoader(
-            train_dataset,
-            batch_size=batch_size,
-            shuffle=True,
-            num_workers=4
-        )
-        self.val_loader = DataLoader(
-            val_dataset,
-            batch_size=batch_size,
-            num_workers=4
-        )
-        self.test_loader = DataLoader(
-            test_dataset,
-            batch_size=batch_size,
-            num_workers=4
-        )
-        
-        self.unique_models = dataset.unique_models
-
-    def train_epoch(self):
-        self.model.train()
-        total_loss = 0
-        predictions = []
-        actual_labels = []
-        
-        progress_bar = tqdm(self.train_loader, desc="Training")
-        for batch_idx, batch in enumerate(progress_bar):
-            try:
-                images, texts, labels = batch
-                images = images.to(self.device)
-                labels = labels.to(self.device)
-                
-                # Move text tensors to device
-                texts = {k: v.squeeze(1).to(self.device) for k, v in texts.items()}
-                
-                self.optimizer.zero_grad()
-                outputs = self.model(images, texts)
-                loss = self.criterion(outputs, labels)
-                
-                loss.backward()
-                # Gradient clipping
-                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
-                self.optimizer.step()
-                
-                total_loss += loss.item()
-                
-                _, preds = torch.max(outputs, 1)
-                predictions.extend(preds.cpu().numpy())
-                actual_labels.extend(labels.cpu().numpy())
-                
-                # Update progress bar
-                progress_bar.set_postfix({
-                    'loss': f'{loss.item():.4f}',
-                    'avg_loss': f'{total_loss/(batch_idx+1):.4f}'
-                })
-                
-            except Exception as e:
-                logging.error(f"Error in batch {batch_idx}: {str(e)}")
-                continue
-        
-        return total_loss / len(self.train_loader), predictions, actual_labels
-
-    def evaluate(self, loader, mode="Validation"):
-        self.model.eval()
-        total_loss = 0
-        predictions = []
-        actual_labels = []
-        
-        with torch.no_grad():
-            progress_bar = tqdm(loader, desc=mode)
-            for batch_idx, batch in enumerate(progress_bar):
-                try:
-                    images, texts, labels = batch
-                    images = images.to(self.device)
-                    labels = labels.to(self.device)
-                    texts = {k: v.squeeze(1).to(self.device) for k, v in texts.items()}
-                    
-                    outputs = self.model(images, texts)
-                    loss = self.criterion(outputs, labels)
-                    
-                    total_loss += loss.item()
-                    
-                    _, preds = torch.max(outputs, 1)
-                    predictions.extend(preds.cpu().numpy())
-                    actual_labels.extend(labels.cpu().numpy())
-                    
-                    progress_bar.set_postfix({
-                        'loss': f'{loss.item():.4f}',
-                        'avg_loss': f'{total_loss/(batch_idx+1):.4f}'
-                    })
-                    
-                except Exception as e:
-                    logging.error(f"Error in {mode} batch {batch_idx}: {str(e)}")
-                    continue
-        
-        return total_loss / len(loader), predictions, actual_labels
-
-    def plot_confusion_matrix(self, y_true, y_pred, title):
-        cm = confusion_matrix(y_true, y_pred)
-        plt.figure(figsize=(15, 15))
-        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
-        plt.title(title)
-        plt.ylabel('True Label')
-        plt.xlabel('Predicted Label')
-        
-        # Save plot
-        filename = f'output/{title.lower().replace(" ", "_")}.png'
-        plt.savefig(filename)
-        plt.close()
-        logging.info(f"Saved confusion matrix to {filename}")
-
-    def generate_evaluation_report(self, y_true, y_pred, title):
-        report = classification_report(
-            y_true,
-            y_pred,
-            target_names=self.unique_models,
-            output_dict=True
-        )
-        df_report = pd.DataFrame(report).transpose()
-        
-        # Save report
-        filename = f'output/{title.lower().replace(" ", "_")}_report.csv'
-        df_report.to_csv(filename)
-        logging.info(f"Saved classification report to {filename}")
-        
-        accuracy = accuracy_score(y_true, y_pred)
-        
-        logging.info(f"\n{title} Results:")
-        logging.info(f"Accuracy: {accuracy:.4f}")
-        logging.info("\nClassification Report:")
-        logging.info("\n" + classification_report(y_true, y_pred, target_names=self.unique_models))
-        
-        return accuracy, df_report
-
-    def train_and_evaluate(self, num_epochs=5):
-        best_val_loss = float('inf')
-        train_accuracies = []
-        val_accuracies = []
-        train_losses = []
-        val_losses = []
-        
-        logging.info(f"Starting training for {num_epochs} epochs...")
-        
-        for epoch in range(num_epochs):
-            logging.info(f"\nEpoch {epoch+1}/{num_epochs}")
-            
-            # Training
-            train_loss, train_preds, train_labels = self.train_epoch()
-            train_accuracy, _ = self.generate_evaluation_report(
-                train_labels,
-                train_preds,
-                f"Training_Epoch_{epoch+1}"
-            )
-            self.plot_confusion_matrix(
-                train_labels,
-                train_preds,
-                f"Training_Confusion_Matrix_Epoch_{epoch+1}"
-            )
-            
-            # Validation
-            val_loss, val_preds, val_labels = self.evaluate(self.val_loader)
-            val_accuracy, _ = self.generate_evaluation_report(
-                val_labels,
-                val_preds,
-                f"Validation_Epoch_{epoch+1}"
-            )
-            self.plot_confusion_matrix(
-                val_labels,
-                val_preds,
-                f"Validation_Confusion_Matrix_Epoch_{epoch+1}"
-            )
-            
-            # Update learning rate scheduler
-            self.scheduler.step(val_loss)
-            
-            train_accuracies.append(train_accuracy)
-            val_accuracies.append(val_accuracy)
-            train_losses.append(train_loss)
-            val_losses.append(val_loss)
-            
-            logging.info(f"\nTraining Loss: {train_loss:.4f}")
-            logging.info(f"Validation Loss: {val_loss:.4f}")
-            
-            # Save best model
-            if val_loss < best_val_loss:
-                best_val_loss = val_loss
-                torch.save({
-                    'epoch': epoch,
-                    'model_state_dict': self.model.state_dict(),
-                    'optimizer_state_dict': self.optimizer.state_dict(),
-                    'val_loss': val_loss,
-                }, 'output/best_model.pth')
-                logging.info(f"Saved new best model with validation loss: {val_loss:.4f}")
-        
-        # Plot training history
-        plt.figure(figsize=(12, 4))
-        
-        # Plot accuracies
-        plt.subplot(1, 2, 1)
-        plt.plot(train_accuracies, label='Training Accuracy')
-        plt.plot(val_accuracies, label='Validation Accuracy')
-        plt.title('Model Accuracy over Epochs')
-        plt.xlabel('Epoch')
-        plt.ylabel('Accuracy')
-        plt.legend()
-        
-        # Plot losses
-        plt.subplot(1, 2, 2)
-        plt.plot(train_losses, label='Training Loss')
-        plt.plot(val_losses, label='Validation Loss')
-        plt.title('Model Loss over Epochs')
-        plt.xlabel('Epoch')
-        plt.ylabel('Loss')
-        plt.legend()
-        
-        plt.tight_layout()
-        plt.savefig('output/training_history.png')
-        plt.close()
-        
-        # Final test evaluation using best model
-        logging.info("\nPerforming final evaluation on test set...")
-        checkpoint = torch.load('output/best_model.pth')
-        self.model.load_state_dict(checkpoint['model_state_dict'])
-        test_loss, test_preds, test_labels = self.evaluate(self.test_loader, "Test")
-        self.generate_evaluation_report(test_labels, test_preds, "Final_Test")
-        self.plot_confusion_matrix(test_labels, test_preds, "Final_Test_Confusion_Matrix")
+# Instantiate model
+model = CombinedModel()
 
+# Define predict function
 def predict(image):
     model.eval()
     with torch.no_grad():
@@ -387,57 +96,14 @@ def predict(image):
         recommended_models = [dataset['Model'][i] for i in indices[0]]
     return recommended_models
 
-def main():
-    try:
-        # Create dataset
-        logging.info("Creating custom dataset...")
-        custom_dataset = CustomDataset(dataset)
-        
-        # Create model
-        logging.info("Initializing model...")
-        model = CombinedModel(num_classes=len(custom_dataset.unique_models))
-        
-        # Create trainer/evaluator
-        logging.info("Setting up trainer/evaluator...")
-        trainer = ModelTrainerEvaluator(
-            model=model,
-            dataset=custom_dataset,
-            batch_size=32,
-            learning_rate=0.001
-        )
-        
-        # Train and evaluate
-        logging.info("Starting training process...")
-        trainer.train_and_evaluate(num_epochs=5)
-        
-        # Create Gradio interface
-        logging.info("Setting up Gradio interface...")
-        interface = gr.Interface(
-            fn=predict,
-            inputs=gr.Image(type="pil"),
-            outputs=gr.Textbox(label="Recommended Models"),
-            title="AI Image Model Recommender",
-            description="Upload an AI-generated image to receive model recommendations.",
-            examples=[
-                ["example_image1.jpg"],
-                ["example_image2.jpg"]
-            ],
-            analytics_enabled=False
-        )
-        
-        # Launch the interface
-        logging.info("Launching Gradio interface...")
-        interface.launch(share=True)
-        
-    except Exception as e:
-        logging.error(f"Error in main function: {str(e)}")
-        raise
+# Set up Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Textbox(label="Recommended Models"),
+    title="AI Image Model Recommender",
+    description="Upload an AI-generated image to receive model recommendations."
+)
 
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        logging.info("Process interrupted by user")
-    except Exception as e:
-        logging.error(f"Fatal error: {str(e)}")
-        raise
+# Launch the app
+interface.launch()