Update app.py

app.py

@@ -1,164 +1,205 @@
-import torch
-import torch.nn as nn
-from torchvision import transforms, models
+import os
+import requests
+from tqdm import tqdm
+from datasets import load_dataset
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
+from tensorflow.keras.preprocessing import image
+from tensorflow.keras.layers import Dense, Input, Concatenate, Embedding, Flatten
+from tensorflow.keras.models import Model
+from tensorflow.keras.preprocessing.text import Tokenizer
+from tensorflow.keras.preprocessing.sequence import pad_sequences
+from sklearn.preprocessing import LabelEncoder
+import joblib
+from PIL import UnidentifiedImageError, Image
 import gradio as gr
-from PIL import Image
 
-# Model Architecture (sama seperti sebelumnya)
-class ModelRecommender(nn.Module):
-    def __init__(self, num_models, text_embedding_dim=768):
-        super(ModelRecommender, self).__init__()
-        
-        # CNN for image processing
-        self.cnn = models.resnet18(pretrained=True)
-        self.cnn.fc = nn.Linear(512, 256)
-        
-        # MLP for text processing
-        self.text_mlp = nn.Sequential(
-            nn.Linear(text_embedding_dim, 512),
-            nn.ReLU(),
-            nn.Linear(512, 256),
-            nn.ReLU()
-        )
-        
-        # Combined layers
-        self.combined = nn.Sequential(
-            nn.Linear(512, 256),
-            nn.ReLU(),
-            nn.Dropout(0.5),
-            nn.Linear(256, num_models)
-        )
-        
-    def forward(self, image, text_features):
-        # Process image
-        img_features = self.cnn(image)
-        
-        # Process text
-        text_features = self.text_mlp(text_features)
-        
-        # Combine features
-        combined = torch.cat((img_features, text_features), dim=1)
-        
-        # Final prediction
-        output = self.combined(combined)
-        return output
+# Constants
+MAX_TEXT_LENGTH = 200
+EMBEDDING_DIM = 100
+IMAGE_SIZE = 224
+BATCH_SIZE = 32
+
+def load_and_preprocess_data(subset_size=2700):
+    # Load dataset
+    dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
+    dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
+    
+    # Filter out NSFW content
+    dataset_subset = dataset_subset.filter(lambda x: not x['nsfw'])
+    
+    return dataset_subset
 
+def process_text_data(dataset_subset):
+    # Combine prompt and negative prompt
+    text_data = [f"{sample['prompt']} {sample['negativePrompt']}" for sample in dataset_subset]
+    
+    # Tokenize text
+    tokenizer = Tokenizer()
+    tokenizer.fit_on_texts(text_data)
+    sequences = tokenizer.texts_to_sequences(text_data)
+    text_data_padded = pad_sequences(sequences, maxlen=MAX_TEXT_LENGTH)
+    
+    return text_data_padded, tokenizer
+
+def process_image_data(dataset_subset):
+    image_dir = 'civitai_images'
+    os.makedirs(image_dir, exist_ok=True)
+    
+    image_data = []
+    valid_indices = []
+    
+    for idx, sample in enumerate(tqdm(dataset_subset)):
+        img_url = sample['url']
+        img_path = os.path.join(image_dir, os.path.basename(img_url))
         
+        try:
+            # Download and save image
+            response = requests.get(img_url)
+            response.raise_for_status()
+            
+            if 'image' not in response.headers['Content-Type']:
+                continue
+                
+            with open(img_path, 'wb') as f:
+                f.write(response.content)
+            
+            # Load and preprocess image
+            img = image.load_img(img_path, target_size=(IMAGE_SIZE, IMAGE_SIZE))
+            img_array = image.img_to_array(img)
+            img_array = preprocess_input(img_array)
+            
+            image_data.append(img_array)
+            valid_indices.append(idx)
+            
+        except Exception as e:
+            print(f"Error processing image {img_url}: {e}")
+            continue
+            
+    return np.array(image_data), valid_indices
 
-# Load model dan dataset info
-def load_model():
-    # Load dataset info
-    dataset_info = torch.load('dataset_info.pth')
-    model_names = dataset_info['model_names']
-    
-    # Initialize model
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = ModelRecommender(len(model_names))
-    
-    # Load model weights
-    checkpoint = torch.load('sd_recommender_model.pth', map_location=device)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.to(device)
-    model.eval()
-    
-    return model, model_names, device
+def create_multimodal_model(num_words, num_classes):
+    # Image input branch (CNN)
+    image_input = Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3))
+    cnn_base = ResNet50(weights='imagenet', include_top=False, pooling='avg')
+    cnn_features = cnn_base(image_input)
+    
+    # Text input branch (MLP)
+    text_input = Input(shape=(MAX_TEXT_LENGTH,))
+    embedding_layer = Embedding(num_words, EMBEDDING_DIM)(text_input)
+    flatten_text = Flatten()(embedding_layer)
+    text_features = Dense(256, activation='relu')(flatten_text)
+    
+    # Combine features
+    combined = Concatenate()([cnn_features, text_features])
+    
+    # Fully connected layers
+    x = Dense(512, activation='relu')(combined)
+    x = Dense(256, activation='relu')(x)
+    output = Dense(num_classes, activation='softmax')(x)
+    
+    model = Model(inputs=[image_input, text_input], outputs=output)
+    return model
 
-def calculate_euclidean_distance(features1, features2):
-    return np.linalg.norm(features1 - features2)
+def train_model():
+    # Load and preprocess data
+    dataset_subset = load_and_preprocess_data()
+    
+    # Process text data
+    text_data_padded, tokenizer = process_text_data(dataset_subset)
+    
+    # Process image data
+    image_data, valid_indices = process_image_data(dataset_subset)
+    
+    # Get valid text data and labels
+    text_data_padded = text_data_padded[valid_indices]
+    model_names = [dataset_subset[i]['Model'] for i in valid_indices]
+    
+    # Encode labels
+    label_encoder = LabelEncoder()
+    encoded_labels = label_encoder.fit_transform(model_names)
+    
+    # Create and compile model
+    model = create_multimodal_model(
+        num_words=len(tokenizer.word_index) + 1,
+        num_classes=len(label_encoder.classes_)
+    )
+    
+    model.compile(
+        optimizer='adam',
+        loss='sparse_categorical_crossentropy',
+        metrics=['accuracy']
+    )
+    
+    # Train model
+    history = model.fit(
+        [image_data, text_data_padded],
+        encoded_labels,
+        batch_size=BATCH_SIZE,
+        epochs=10,
+        validation_split=0.2
+    )
+    
+    # Save models and encoders
+    model.save('multimodal_model')
+    joblib.dump(tokenizer, 'tokenizer.pkl')
+    joblib.dump(label_encoder, 'label_encoder.pkl')
+    
+    return model, tokenizer, label_encoder
 
-def predict_image(image):
-    if not hasattr(predict_image, "model"):
-        predict_image.model, predict_image.model_names, predict_image.device = load_model()
-    
-    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])
-    ])
-    
-    image_tensor = transform(image).unsqueeze(0).to(predict_image.device)
-    dummy_text_features = torch.zeros(1, 768).to(predict_image.device)
-    
-    # Get image features
-    with torch.no_grad():
-        img_features = predict_image.model.cnn(image_tensor).cpu().numpy()
-        outputs = predict_image.model(image_tensor, dummy_text_features)
-        top5_prob, top5_indices = torch.topk(outputs, 5)
-    
-    # Create HTML gallery
-    html_output = """
-    <style>
-    .model-gallery {
-        display: grid;
-        grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
-        gap: 20px;
-        padding: 20px;
-    }
-    .model-card {
-        border: 1px solid #ddd;
-        border-radius: 8px;
-        overflow: hidden;
-        background: white;
-        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
-    }
-    .model-img {
-        width: 100%;
-        height: 200px;
-        object-fit: cover;
-    }
-    .model-info {
-        padding: 15px;
-    }
-    .model-name {
-        color: #2563eb;
-        text-decoration: none;
-        font-weight: bold;
-        font-size: 1.1em;
-        margin-bottom: 8px;
-        display: block;
-    }
-    .model-name:hover {
-        text-decoration: underline;
-    }
-    .distance {
-        color: #666;
-        font-size: 0.9em;
-    }
-    </style>
-    <div class="model-gallery">
-    """
-    
-    # Generate cards for each model
-    for idx, (score, model_idx) in enumerate(zip(top5_prob[0], top5_indices[0])):
-        model_name = predict_image.model_names[model_idx.item()]
-        distance = calculate_euclidean_distance(img_features[0], 
-                                             torch.randn(512).numpy())  # Placeholder for actual features
-        
-        civitai_url = f"https://civitai.com/search/models?sortBy=models_v9&query={model_name}"
-        
-        html_output += f"""
-        <div class="model-card">
-            <img class="model-img" src="data:image/svg+xml,<svg xmlns='http://www.w3.org/2000/svg' width='250' height='200' viewBox='0 0 250 200'><rect width='100%' height='100%' fill='%23f0f0f0'/><text x='50%' y='50%' dominant-baseline='middle' text-anchor='middle' font-family='Arial' font-size='16' fill='%23666'>Model Preview</text></svg>" alt="{model_name}">
-            <div class="model-info">
-                <a href="{civitai_url}" target="_blank" class="model-name">{model_name}</a>
-                <div class="distance">Euclidean Distance: {distance:.4f}</div>
-            </div>
-        </div>
-        """
-    
-    html_output += "</div>"
-    
-    return html_output
+def get_recommendations(image_input, text_input, model, tokenizer, label_encoder, top_k=5):
+    # Preprocess image
+    img_array = image.img_to_array(image_input)
+    img_array = tf.image.resize(img_array, (IMAGE_SIZE, IMAGE_SIZE))
+    img_array = preprocess_input(img_array)
+    img_array = np.expand_dims(img_array, axis=0)
+    
+    # Preprocess text
+    text_sequence = tokenizer.texts_to_sequences([text_input])
+    text_padded = pad_sequences(text_sequence, maxlen=MAX_TEXT_LENGTH)
+    
+    # Get predictions
+    predictions = model.predict([img_array, text_padded])
+    top_indices = np.argsort(predictions[0])[-top_k:][::-1]
+    
+    # Get recommended model names and confidence scores
+    recommendations = [
+        (label_encoder.inverse_transform([idx])[0], predictions[0][idx])
+        for idx in top_indices
+    ]
+    
+    return recommendations
 
-# Gradio Interface
-demo = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.HTML(),
-    title="Stable Diffusion Model Recommender",
-    description="Upload an AI-generated image to get model recommendations",
-    examples=[["example1.jpg"], ["example2.jpg"]]
-)
+# Gradio interface
+def create_gradio_interface():
+    # Load saved models
+    model = tf.keras.models.load_model('multimodal_model')
+    tokenizer = joblib.load('tokenizer.pkl')
+    label_encoder = joblib.load('label_encoder.pkl')
+    
+    def predict(img, text):
+        recommendations = get_recommendations(img, text, model, tokenizer, label_encoder)
+        return "\n".join([f"Model: {name}, Confidence: {conf:.2f}" for name, conf in recommendations])
+    
+    interface = gr.Interface(
+        fn=predict,
+        inputs=[
+            gr.Image(type="pil", label="Upload Image"),
+            gr.Textbox(label="Enter Prompt")
+        ],
+        outputs=gr.Textbox(label="Recommended Models"),
+        title="Multimodal Model Recommendation System",
+        description="Upload an image and enter a prompt to get model recommendations"
+    )
+    
+    return interface
 
-demo.launch()
+if __name__ == "__main__":
+    # Train model if not already trained
+    if not os.path.exists('multimodal_model'):
+        model, tokenizer, label_encoder = train_model()
+    
+    # Launch Gradio interface
+    interface = create_gradio_interface()
+    interface.launch()