Update app.py

app.py

@@ -1,192 +1,140 @@
-import torch
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForCausalLM
-from PIL import Image
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from torchvision import models
 import pandas as pd
 from datasets import load_dataset
-from sklearn.metrics.pairwise import cosine_similarity
-import numpy as np
-import warnings
-warnings.filterwarnings('ignore')
-
-# 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
+from torch.utils.data import DataLoader, Dataset
+from sklearn.preprocessing import LabelEncoder
 
-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 dataset
+dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
 
-# Create a dummy image for text-only processing
-DUMMY_IMAGE = Image.new('RGB', (224, 224), color='white')
+# Text preprocessing function
+def preprocess_text(text, max_length=100):
+    # Convert text to lowercase and split into words
+    words = text.lower().split()
+    # Truncate or pad to max_length
+    if len(words) > max_length:
+        words = words[:max_length]
+    else:
+        words.extend([''] * (max_length - len(words)))
+    return words
 
-# Load CivitAI dataset
-print("Loading dataset...")
-dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k", split="train[:1000]")
-df = pd.DataFrame(dataset)
-print("Dataset loaded successfully!")
-
-text_embedding_cache = {}
-
-def get_image_embedding(image):
-    try:
-        inputs = processor(
-            images=image,
-            text="Generate image description",
-            return_tensors="pt",
-            padding=True
-        ).to(device, torch_dtype)
-        
-        # Generate decoder_input_ids with adjusted parameters
-        decoder_input_ids = model.generate(
-            **inputs,
-            max_new_tokens=20,  # Increased from max_length
-            min_length=1,
-            num_beams=1,
-            do_sample=False,
-            pad_token_id=processor.tokenizer.pad_token_id,
-            return_dict_in_generate=True,
-        ).sequences
-        
-        inputs['decoder_input_ids'] = decoder_input_ids
-        
-        with torch.no_grad():
-            outputs = model(**inputs)
-            image_embeddings = outputs.last_hidden_state.mean(dim=1)
-        return image_embeddings.cpu().numpy()
-    except Exception as e:
-        print(f"Error in get_image_embedding: {str(e)}")
-        return None
-
-def get_text_embedding(text):
-    try:
-        if text in text_embedding_cache:
-            return text_embedding_cache[text]
-        
-        # Process text with dummy image
-        inputs = processor(
-            images=DUMMY_IMAGE,
-            text=text,
-            return_tensors="pt",
-            padding=True
-        ).to(device, torch_dtype)
-        
-        # Generate decoder_input_ids with adjusted parameters
-        decoder_input_ids = model.generate(
-            **inputs,
-            max_new_tokens=20,  # Using max_new_tokens instead of max_length
-            min_length=1,
-            num_beams=1,
-            do_sample=False,
-            pad_token_id=processor.tokenizer.pad_token_id,
-            return_dict_in_generate=True,
-        ).sequences
+class CustomDataset(Dataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+        ])
+        self.label_encoder = LabelEncoder()
+        self.labels = self.label_encoder.fit_transform(dataset['Model'])
         
-        inputs['decoder_input_ids'] = decoder_input_ids
+        # Create vocabulary from all prompts
+        self.vocab = set()
+        for item in dataset['prompt']:
+            self.vocab.update(preprocess_text(item))
+        self.vocab = list(self.vocab)
+        self.word_to_idx = {word: idx for idx, word in enumerate(self.vocab)}
         
-        with torch.no_grad():
-            outputs = model(**inputs)
-            text_embeddings = outputs.last_hidden_state.mean(dim=1)
-        
-        embedding = text_embeddings.cpu().numpy()
-        text_embedding_cache[text] = embedding
-        return embedding
-    except Exception as e:
-        print(f"Error in get_text_embedding: {str(e)}")
-        return None
-
-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")
-
-def find_similar_images(uploaded_image, top_k=5):
-    query_embedding = get_image_embedding(uploaded_image)
-    if query_embedding is None:
-        return [], []
+    def __len__(self):
+        return len(self.dataset)
     
-    similarities = []
-    for idx, row in df.iterrows():
-        prompt_embedding = get_text_embedding(row['prompt'])
-        if prompt_embedding is not None:
-            similarity = cosine_similarity(query_embedding, prompt_embedding)[0][0]
-            similarities.append({
-                'similarity': similarity,
-                'model': row['Model'],
-                'prompt': row['prompt']
-            })
+    def text_to_vector(self, text):
+        words = preprocess_text(text)
+        vector = torch.zeros(len(self.vocab))
+        for word in words:
+            if word in self.word_to_idx:
+                vector[self.word_to_idx[word]] += 1
+        return vector
     
-    sorted_results = sorted(similarities, key=lambda x: x['similarity'], reverse=True)
-    top_models = []
-    top_prompts = []
-    seen_models = set()
-    seen_prompts = set()
-    
-    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'])
+    def __getitem__(self, idx):
+        image = self.transform(self.dataset[idx]['image'])
+        text_vector = self.text_to_vector(self.dataset[idx]['prompt'])
+        label = self.labels[idx]
+        return image, text_vector, 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)
         
-        if len(top_prompts) < top_k and result['prompt'] not in seen_prompts:
-            top_prompts.append(result['prompt'])
-            seen_prompts.add(result['prompt'])
+    def forward(self, x):
+        return self.model(x)
+
+# Define MLP for text processing
+class TextMLP(nn.Module):
+    def __init__(self, vocab_size):
+        super(TextMLP, self).__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(vocab_size, 1024),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(512, 512)
+        )
         
-        if len(top_models) == top_k and len(top_prompts) == top_k:
-            break
-    
-    return top_models, top_prompts
+    def forward(self, x):
+        return self.layers(x)
 
-def process_image(input_image):
-    if input_image is None:
-        return "Please upload an image.", "Please upload an image."
-    
-    try:
-        if not isinstance(input_image, Image.Image):
-            input_image = Image.fromarray(input_image)
+# Combined model
+class CombinedModel(nn.Module):
+    def __init__(self, vocab_size):
+        super(CombinedModel, self).__init__()
+        self.image_model = ImageModel()
+        self.text_model = TextMLP(vocab_size)
+        self.fc = nn.Linear(1024, len(dataset['Model'].unique()))
         
-        # Resize image to expected size
-        input_image = input_image.resize((224, 224))
+    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)
+
+# Create dataset instance and model
+custom_dataset = CustomDataset(dataset)
+model = CombinedModel(len(custom_dataset.vocab))
+
+def get_recommendations(image):
+    model.eval()
+    with torch.no_grad():
+        # Process input image
+        transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor()
+        ])
+        image_tensor = transform(image).unsqueeze(0)
         
-        recommended_models, recommended_prompts = find_similar_images(input_image)
+        # Create dummy text vector (since we're only doing image-based recommendations)
+        dummy_text = torch.zeros((1, len(custom_dataset.vocab)))
         
-        if not recommended_models or not recommended_prompts:
-            return "Error processing image.", "Error processing image."
+        # Get model output
+        output = model(image_tensor, dummy_text)
+        _, indices = torch.topk(output, 5)
         
-        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)])
+        # Get recommended images and their information
+        recommendations = []
+        for idx in indices[0]:
+            recommended_image = dataset[idx.item()]['image']
+            model_name = dataset[idx.item()]['Model']
+            recommendations.append((recommended_image, f"{model_name}"))
         
-        return models_text, prompts_text
-    except Exception as e:
-        print(f"Error in process_image: {str(e)}")
-        return "Error processing image.", "Error processing image."
-
-# Pre-compute embeddings when starting the application
-try:
-    precompute_embeddings()
-except Exception as e:
-    print(f"Error in precompute_embeddings: {str(e)}")
+    return recommendations
 
-# 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
+# Set up Gradio interface
+interface = gr.Interface(
+    fn=get_recommendations,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Gallery(label="Recommended Images"),
+    title="Image Recommendation System",
+    description="Upload an image and get similar images with their model names."
 )
 
-# Launch the interface
-iface.launch()
+# Launch the app
+interface.launch()