app.py

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

# 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

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 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 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
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

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()