Update app.py

app.py

@@ -16,12 +16,12 @@ from PIL import UnidentifiedImageError, Image
 import gradio as gr
 
 # Optimized Constants
-MAX_TEXT_LENGTH = 100  # Reduced from 200
-EMBEDDING_DIM = 50    # Reduced from 100
-IMAGE_SIZE = 160      # Reduced from 224
-BATCH_SIZE = 64      # Increased from 32
+MAX_TEXT_LENGTH = 100
+EMBEDDING_DIM = 50
+IMAGE_SIZE = 160
+BATCH_SIZE = 64
 
-def load_and_preprocess_data(subset_size=500):  # Reduced from 2700
+def load_and_preprocess_data(subset_size=500):
     # Load dataset
     dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
     dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
@@ -36,7 +36,7 @@ def process_text_data(dataset_subset):
     text_data = [f"{sample['prompt']} {sample['negativePrompt']}" for sample in dataset_subset]
     
     # Tokenize text
-    tokenizer = Tokenizer(num_words=10000)  # Added limit to vocabulary size
+    tokenizer = Tokenizer(num_words=10000)
     tokenizer.fit_on_texts(text_data)
     sequences = tokenizer.texts_to_sequences(text_data)
     text_data_padded = pad_sequences(sequences, maxlen=MAX_TEXT_LENGTH)
@@ -56,7 +56,7 @@ def process_image_data(dataset_subset):
         
         try:
             # Download and save image
-            response = requests.get(img_url, timeout=5)  # Added timeout
+            response = requests.get(img_url, timeout=5)
             response.raise_for_status()
             
             if 'image' not in response.headers['Content-Type']:
@@ -83,23 +83,23 @@ def create_multimodal_model(num_words, num_classes):
     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 to reduce training time
+    # 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 (Simplified MLP)
+    # 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)  # Reduced from 256
+    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)  # Reduced from 512
+    x = Dense(256, activation='relu')(combined)
     output = Dense(num_classes, activation='softmax')(x)
     
     model = Model(inputs=[image_input, text_input], outputs=output)
@@ -125,7 +125,7 @@ def train_model():
     
     # Create and compile model
     model = create_multimodal_model(
-        num_words=10000,  # Limited vocabulary size
+        num_words=10000,
         num_classes=len(label_encoder.classes_)
     )
     
@@ -135,17 +135,17 @@ def train_model():
         metrics=['accuracy']
     )
     
-    # Train model with reduced epochs
+    # Train model
     history = model.fit(
         [image_data, text_data_padded],
         encoded_labels,
         batch_size=BATCH_SIZE,
-        epochs=3,  # Reduced from 10
+        epochs=3,
         validation_split=0.2
     )
     
-    # Save models and encoders
-    model.save('multimodal_model')
+    # Save models and encoders with correct extensions
+    model.save('multimodal_model.keras')  # Changed from 'multimodal_model'
     joblib.dump(tokenizer, 'tokenizer.pkl')
     joblib.dump(label_encoder, 'label_encoder.pkl')
     
@@ -174,10 +174,9 @@ def get_recommendations(image_input, text_input, model, tokenizer, label_encoder
     
     return recommendations
 
-# Gradio interface
 def create_gradio_interface():
-    # Load saved models
-    model = tf.keras.models.load_model('multimodal_model')
+    # Load saved models with correct path
+    model = tf.keras.models.load_model('multimodal_model.keras')  # Changed from 'multimodal_model'
     tokenizer = joblib.load('tokenizer.pkl')
     label_encoder = joblib.load('label_encoder.pkl')
     
@@ -200,8 +199,12 @@ def create_gradio_interface():
 
 if __name__ == "__main__":
     # Train model if not already trained
-    if not os.path.exists('multimodal_model'):
+    if not os.path.exists('multimodal_model.keras'):  # Changed from 'multimodal_model'
+        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()