Create app.py

app.py

@@ -0,0 +1,228 @@
+# app.py
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+from transformers import AutoModel, AutoProcessor
+from torch import nn
+import torch.nn.functional as F
+from datasets import load_dataset
+from torch.utils.data import Dataset, DataLoader
+import os
+from tqdm import tqdm
+
+class SDDataset(Dataset):
+    def __init__(self, dataset, processor, model_to_idx, token_to_idx, max_samples=5000):
+        self.dataset = dataset.select(range(min(max_samples, len(dataset))))
+        self.processor = processor
+        self.model_to_idx = model_to_idx
+        self.token_to_idx = token_to_idx
+        
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, idx):
+        item = self.dataset[idx]
+        
+        # Process image
+        image = Image.open(item['image'])
+        image_inputs = self.processor(images=image, return_tensors="pt")
+        
+        # Create model label
+        model_label = torch.zeros(len(self.model_to_idx))
+        model_label[self.model_to_idx[item['model_name']]] = 1
+        
+        # Create prompt label (multi-hot encoding)
+        prompt_label = torch.zeros(len(self.token_to_idx))
+        for token in item['prompt'].split():
+            if token in self.token_to_idx:
+                prompt_label[self.token_to_idx[token]] = 1
+                
+        return image_inputs, model_label, prompt_label
+
+class SDRecommenderModel(nn.Module):
+    def __init__(self, florence_model, num_models, vocab_size):
+        super().__init__()
+        self.florence = florence_model
+        self.model_head = nn.Linear(florence_model.config.hidden_size, num_models)
+        self.prompt_head = nn.Linear(florence_model.config.hidden_size, vocab_size)
+        
+    def forward(self, image_features):
+        # Get Florence embeddings
+        features = self.florence.get_image_features(image_features)
+        
+        # Generate model and prompt recommendations
+        model_logits = self.model_head(features)
+        prompt_logits = self.prompt_head(features)
+        
+        return model_logits, prompt_logits
+
+class SDRecommender:
+    def __init__(self, max_samples=1000):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+        
+        # Load Florence model and processor
+        self.processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large")
+        self.florence = AutoModel.from_pretrained("microsoft/Florence-2-large")
+        
+        # Load dataset
+        print("Loading dataset...")
+        self.dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k", split="train")
+        self.dataset = self.dataset.select(range(min(max_samples, len(self.dataset))))
+        print(f"Using {len(self.dataset)} samples from dataset")
+        
+        # Create vocabularies for models and tokens
+        self.model_to_idx = self._create_model_vocab()
+        self.token_to_idx = self._create_prompt_vocab()
+        
+        # Initialize the recommendation model
+        self.model = SDRecommenderModel(
+            self.florence,
+            len(self.model_to_idx),
+            len(self.token_to_idx)
+        ).to(self.device)
+        
+        # Load trained weights if available
+        if os.path.exists("recommender_model.pth"):
+            self.model.load_state_dict(torch.load("recommender_model.pth"))
+            print("Loaded trained model weights")
+        self.model.eval()
+        
+    def _create_model_vocab(self):
+        print("Creating model vocabulary...")
+        models = set()
+        for item in self.dataset:
+            models.add(item["model_name"])
+        return {model: idx for idx, model in enumerate(sorted(models))}
+    
+    def _create_prompt_vocab(self):
+        print("Creating prompt vocabulary...")
+        tokens = set()
+        for item in self.dataset:
+            for token in item["prompt"].split():
+                tokens.add(token)
+        return {token: idx for idx, token in enumerate(sorted(tokens))}
+    
+    def train(self, num_epochs=5, batch_size=8, learning_rate=1e-4):
+        print("Starting training...")
+        
+        # Create dataset and dataloader
+        train_dataset = SDDataset(
+            self.dataset,
+            self.processor,
+            self.model_to_idx,
+            self.token_to_idx
+        )
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=2
+        )
+        
+        # Setup optimizer
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=learning_rate)
+        
+        # Training loop
+        self.model.train()
+        for epoch in range(num_epochs):
+            total_loss = 0
+            progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")
+            
+            for batch_idx, (images, model_labels, prompt_labels) in enumerate(progress_bar):
+                # Move everything to device
+                images = images.to(self.device)
+                model_labels = model_labels.to(self.device)
+                prompt_labels = prompt_labels.to(self.device)
+                
+                # Forward pass
+                model_logits, prompt_logits = self.model(images)
+                
+                # Calculate loss
+                model_loss = F.cross_entropy(model_logits, model_labels)
+                prompt_loss = F.binary_cross_entropy_with_logits(prompt_logits, prompt_labels)
+                loss = model_loss + prompt_loss
+                
+                # Backward pass
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                
+                # Update progress
+                total_loss += loss.item()
+                progress_bar.set_postfix({"loss": total_loss / (batch_idx + 1)})
+        
+        # Save trained model
+        torch.save(self.model.state_dict(), "recommender_model.pth")
+        print("Training completed and model saved")
+    
+    def get_recommendations(self, image):
+        # Convert uploaded image to PIL if needed
+        if not isinstance(image, Image.Image):
+            image = Image.open(image)
+        
+        # Process image
+        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+        
+        # Get model predictions
+        self.model.eval()
+        with torch.no_grad():
+            model_logits, prompt_logits = self.model(inputs)
+            
+            # Get top 5 model recommendations
+            model_probs = F.softmax(model_logits, dim=-1)
+            top_models = torch.topk(model_probs, k=5)
+            model_recommendations = [
+                (list(self.model_to_idx.keys())[idx.item()], prob.item())
+                for prob, idx in zip(top_models.values[0], top_models.indices[0])
+            ]
+            
+            # Get top tokens for prompt recommendations
+            prompt_probs = F.softmax(prompt_logits, dim=-1)
+            top_tokens = torch.topk(prompt_probs, k=20)
+            recommended_tokens = [
+                list(self.token_to_idx.keys())[idx.item()]
+                for idx in top_tokens.indices[0]
+            ]
+            
+            # Create 5 prompt combinations
+            prompt_recommendations = [
+                " ".join(np.random.choice(recommended_tokens, size=8, replace=False))
+                for _ in range(5)
+            ]
+            
+        return (
+            "\n".join(f"{model} (confidence: {conf:.2f})" for model, conf in model_recommendations),
+            "\n".join(prompt_recommendations)
+        )
+
+# Create Gradio interface
+def create_interface():
+    recommender = SDRecommender(max_samples=5000)
+    
+    # Train the model if no trained weights exist
+    if not os.path.exists("recommender_model.pth"):
+        recommender.train()
+    
+    def process_image(image):
+        model_recs, prompt_recs = recommender.get_recommendations(image)
+        return model_recs, prompt_recs
+    
+    interface = gr.Interface(
+        fn=process_image,
+        inputs=gr.Image(type="pil"),
+        outputs=[
+            gr.Textbox(label="Recommended Models"),
+            gr.Textbox(label="Recommended Prompts")
+        ],
+        title="Stable Diffusion Model & Prompt Recommender",
+        description="Upload an AI-generated image to get model and prompt recommendations",
+    )
+    
+    return interface
+
+# Launch the interface
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch()