Update app.py

app.py

@@ -1,239 +1,128 @@
-# app.py
-import gradio as gr
 import torch
-import numpy as np
+import gradio as gr
+from transformers import AutoProcessor, AutoModelForCausalLM
 from PIL import Image
-from torch import nn
-import torch.nn.functional as F
+import pandas as pd
 from datasets import load_dataset
-from torch.utils.data import Dataset, DataLoader
-import os
-from tqdm import tqdm
-from transformers import AutoProcessor, AutoModelForCausalLM
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
 
-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)
+# Load Florence-2 model and processor
+model_name = "microsoft/Florence-2-base"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch_dtype,
+    trust_remote_code=True
+).to(device)
+processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+
+# Load CivitAI dataset (limited to 1000 samples)
+dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k", split="train[:1000]")
+df = pd.DataFrame(dataset)
+
+# Create cache for embeddings to improve performance
+text_embedding_cache = {}
+
+def get_image_embedding(image):
+    inputs = processor(images=image, return_tensors="pt").to(device, torch_dtype)
+    with torch.no_grad():
+        outputs = model.get_image_features(**inputs)
+    return outputs.cpu().numpy()
+
+def get_text_embedding(text):
+    if text in text_embedding_cache:
+        return text_embedding_cache[text]
     
-    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
+    inputs = processor(text=text, return_tensors="pt").to(device, torch_dtype)
+    with torch.no_grad():
+        outputs = model.get_text_features(**inputs)
+    
+    embedding = outputs.cpu().numpy()
+    text_embedding_cache[text] = embedding
+    return embedding
 
-class SDRecommenderModel(nn.Module):
-    def __init__(self, florence_model, num_models, vocab_size):
-        super().__init__()
-        self.florence = florence_model
-        hidden_size = 1024  # Florence-2-large hidden size
-        self.model_head = nn.Linear(hidden_size, num_models)
-        self.prompt_head = nn.Linear(hidden_size, vocab_size)
-        
-    def forward(self, pixel_values):
-        # Get Florence embeddings
-        outputs = self.florence(pixel_values=pixel_values, output_hidden_states=True)
-        features = outputs.hidden_states[-1].mean(dim=1)  # Use mean pooling of last hidden state
-        
-        # Generate model and prompt recommendations
-        model_logits = self.model_head(features)
-        prompt_logits = self.prompt_head(features)
-        
-        return model_logits, prompt_logits
+# Pre-compute text embeddings for all prompts in the dataset
+def precompute_embeddings():
+    print("Pre-computing text embeddings...")
+    for idx, row in df.iterrows():
+        if row['prompt'] not in text_embedding_cache:
+            _ = get_text_embedding(row['prompt'])
+        if idx % 100 == 0:
+            print(f"Processed {idx}/1000 embeddings")
+    print("Finished pre-computing embeddings")
 
-class SDRecommender:
-    def __init__(self, max_samples=500):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {self.device}")
-        
-        # Load Florence model and processor
-        print("Loading Florence model and processor...")
-        self.processor = AutoProcessor.from_pretrained(
-            "microsoft/Florence-2-large",
-            trust_remote_code=True
-        )
-        self.florence = AutoModelForCausalLM.from_pretrained(
-            "microsoft/Florence-2-large",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            trust_remote_code=True
-        ).to(self.device)
-        
-        # 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", map_location=self.device))
-            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 find_similar_images(uploaded_image, top_k=5):
+    # Get embedding for uploaded image
+    query_embedding = get_image_embedding(uploaded_image)
     
-    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))}
+    # Calculate similarities with dataset
+    similarities = []
+    for idx, row in df.iterrows():
+        prompt_embedding = get_text_embedding(row['prompt'])
+        similarity = cosine_similarity(query_embedding, prompt_embedding)[0][0]
+        similarities.append({
+            'similarity': similarity,
+            'model': row['Model'],
+            'prompt': row['prompt']
+        })
     
-    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 = {k: v.to(self.device) for k, v in images.items()}
-                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")
+    # Sort by similarity and get top k results
+    sorted_results = sorted(similarities, key=lambda x: x['similarity'], reverse=True)
+    top_models = []
+    top_prompts = []
+    seen_models = set()
+    seen_prompts = set()
     
-    def get_recommendations(self, image):
-        # Convert uploaded image to PIL if needed
-        if not isinstance(image, Image.Image):
-            image = Image.open(image)
+    for result in sorted_results:
+        if len(top_models) < top_k and result['model'] not in seen_models:
+            top_models.append(result['model'])
+            seen_models.add(result['model'])
         
-        # Process image
-        inputs = self.processor(images=image, return_tensors="pt")
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        if len(top_prompts) < top_k and result['prompt'] not in seen_prompts:
+            top_prompts.append(result['prompt'])
+            seen_prompts.add(result['prompt'])
         
-        # 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)
-        )
+        if len(top_models) == top_k and len(top_prompts) == top_k:
+            break
+    
+    return top_models, top_prompts
 
-# Create Gradio interface
-def create_interface():
-    recommender = SDRecommender(max_samples=5000)
+def process_image(input_image):
+    if input_image is None:
+        return "Please upload an image.", "Please upload an image."
     
-    # Train the model if no trained weights exist
-    if not os.path.exists("recommender_model.pth"):
-        recommender.train()
+    # Convert to PIL Image if needed
+    if not isinstance(input_image, Image.Image):
+        input_image = Image.fromarray(input_image)
     
-    def process_image(image):
-        model_recs, prompt_recs = recommender.get_recommendations(image)
-        return model_recs, prompt_recs
+    # Get recommendations
+    recommended_models, recommended_prompts = find_similar_images(input_image)
     
-    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",
-    )
+    # Format output
+    models_text = "Recommended Models:\n" + "\n".join([f"{i+1}. {model}" for i, model in enumerate(recommended_models)])
+    prompts_text = "Recommended Prompts:\n" + "\n".join([f"{i+1}. {prompt}" for i, prompt in enumerate(recommended_prompts)])
     
-    return interface
+    return models_text, prompts_text
+
+# Pre-compute embeddings when starting the application
+precompute_embeddings()
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil", label="Upload AI-generated image"),
+    outputs=[
+        gr.Textbox(label="Recommended Models", lines=6),
+        gr.Textbox(label="Recommended Prompts", lines=6)
+    ],
+    title="AI Image Model & Prompt Recommender",
+    description="Upload an AI-generated image to get recommendations for Stable Diffusion models and prompts.",
+    examples=[],
+    cache_examples=False
+)
 
 # Launch the interface
-if __name__ == "__main__":
-    interface = create_interface()
-    interface.launch()
+iface.launch()