Update app.py

app.py

@@ -21,7 +21,10 @@ EMBEDDING_DIM = 50
 IMAGE_SIZE = 160
 BATCH_SIZE = 64
 
-def load_and_preprocess_data(subset_size=10000):
+# Store model examples
+model_examples = {}
+
+def load_and_preprocess_data(subset_size=20000):
     # Load dataset
     dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
     dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
@@ -29,13 +32,17 @@ def load_and_preprocess_data(subset_size=10000):
     # Filter out NSFW content
     dataset_subset = dataset_subset.filter(lambda x: not x['nsfw'])
     
+    # Store example images for each model
+    for item in dataset_subset:
+        if item['Model'] not in model_examples:
+            model_examples[item['Model']] = item['url']
+    
     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]
+    # Combine prompt and negative prompt without user input
+    text_data = ["default prompt" for _ in dataset_subset]
     
-    # Tokenize text
     tokenizer = Tokenizer(num_words=10000)
     tokenizer.fit_on_texts(text_data)
     sequences = tokenizer.texts_to_sequences(text_data)
@@ -43,6 +50,14 @@ def process_text_data(dataset_subset):
     
     return text_data_padded, tokenizer
 
+def download_image(url):
+    try:
+        response = requests.get(url, timeout=5)
+        response.raise_for_status()
+        return Image.open(requests.get(url, stream=True).raw)
+    except:
+        return None
+
 def process_image_data(dataset_subset):
     image_dir = 'civitai_images'
     os.makedirs(image_dir, exist_ok=True)
@@ -55,7 +70,6 @@ def process_image_data(dataset_subset):
         img_path = os.path.join(image_dir, os.path.basename(img_url))
         
         try:
-            # Download and save image
             response = requests.get(img_url, timeout=5)
             response.raise_for_status()
             
@@ -65,7 +79,6 @@ def process_image_data(dataset_subset):
             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)
@@ -79,26 +92,21 @@ def process_image_data(dataset_subset):
     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')
     
-    # Freeze most of the ResNet50 layers
     for layer in cnn_base.layers[:-10]:
         layer.trainable = False
         
     cnn_features = cnn_base(image_input)
     
-    # Text input branch
     text_input = Input(shape=(MAX_TEXT_LENGTH,))
     embedding_layer = Embedding(num_words, EMBEDDING_DIM)(text_input)
     flatten_text = Flatten()(embedding_layer)
     text_features = Dense(128, activation='relu')(flatten_text)
     
-    # Combine features
     combined = Concatenate()([cnn_features, text_features])
     
-    # Simplified fully connected layers
     x = Dense(256, activation='relu')(combined)
     output = Dense(num_classes, activation='softmax')(x)
     
@@ -106,24 +114,18 @@ def create_multimodal_model(num_words, num_classes):
     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=10000,
         num_classes=len(label_encoder.classes_)
@@ -135,7 +137,6 @@ def train_model():
         metrics=['accuracy']
     )
     
-    # Train model
     history = model.fit(
         [image_data, text_data_padded],
         encoded_labels,
@@ -144,68 +145,75 @@ def train_model():
         validation_split=0.2
     )
     
-    # Save models and encoders with correct extensions
-    model.save('multimodal_model.keras')  # Changed from 'multimodal_model'
+    model.save('multimodal_model.keras')
     joblib.dump(tokenizer, 'tokenizer.pkl')
     joblib.dump(label_encoder, 'label_encoder.pkl')
     
+    # Save model examples
+    joblib.dump(model_examples, 'model_examples.pkl')
+    
     return model, tokenizer, label_encoder
 
-def get_recommendations(image_input, text_input, model, tokenizer, label_encoder, top_k=5):
-    # Preprocess image
+def get_recommendations(image_input, model, tokenizer, label_encoder, top_k=5):
     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])
+    # Use default text input
+    text_sequence = tokenizer.texts_to_sequences(["default prompt"])
     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
-    ]
+    recommendations = []
+    for idx in top_indices:
+        model_name = label_encoder.inverse_transform([idx])[0]
+        confidence = predictions[0][idx]
+        if model_name in model_examples:
+            example_image = download_image(model_examples[model_name])
+            if example_image:
+                recommendations.append((model_name, confidence, example_image))
     
     return recommendations
 
 def create_gradio_interface():
-    # Load saved models with correct path
-    model = tf.keras.models.load_model('multimodal_model.keras')  # Changed from 'multimodal_model'
+    model = tf.keras.models.load_model('multimodal_model.keras')
     tokenizer = joblib.load('tokenizer.pkl')
     label_encoder = joblib.load('label_encoder.pkl')
+    model_examples_data = joblib.load('model_examples.pkl')
+    
+    def predict(img):
+        recommendations = get_recommendations(img, model, tokenizer, label_encoder)
+        result_text = ""
+        result_images = []
+        
+        for model_name, conf, example_img in recommendations:
+            result_text += f"Model: {model_name}\n"
+            result_images.append(example_img)
+        
+        return [result_text] + result_images
     
-    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])
+    outputs = [gr.Textbox(label="Recommended Models")] + [gr.Image(label=f"Example {i+1}") for i in range(5)]
     
     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"
+        inputs=gr.Image(type="pil", label="Upload Image"),
+        outputs=outputs,
+        title="AI Model Recommendation System",
+        description="Upload an image to get model recommendations with examples"
     )
     
     return interface
 
 if __name__ == "__main__":
-    # Train model if not already trained
-    if not os.path.exists('multimodal_model.keras'):  # Changed from 'multimodal_model'
+    if not os.path.exists('multimodal_model.keras'):
         print("Training new model...")
         model, tokenizer, label_encoder = train_model()
         print("Training completed!")
     else:
         print("Loading existing model...")
     
-    # Launch Gradio interface
     interface = create_gradio_interface()
     interface.launch()