Update app.py

app.py

@@ -1,4 +1,4 @@
-import os
+mport os
 import requests
 from tqdm import tqdm
 from datasets import load_dataset
@@ -15,13 +15,13 @@ 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
+# 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
 
-def load_and_preprocess_data(subset_size=2700):
+def load_and_preprocess_data(subset_size=500):  # Reduced from 2700
     # 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()
+    tokenizer = Tokenizer(num_words=10000)  # Added limit to vocabulary size
     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)
+            response = requests.get(img_url, timeout=5)  # Added timeout
             response.raise_for_status()
             
             if 'image' not in response.headers['Content-Type']:
@@ -74,7 +74,6 @@ def process_image_data(dataset_subset):
             valid_indices.append(idx)
             
         except Exception as e:
-            print(f"Error processing image {img_url}: {e}")
             continue
             
     return np.array(image_data), valid_indices
@@ -83,20 +82,24 @@ 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 to reduce training time
+    for layer in cnn_base.layers[:-10]:
+        layer.trainable = False
+        
     cnn_features = cnn_base(image_input)
     
-    # Text input branch (MLP)
+    # Text input branch (Simplified 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)
+    text_features = Dense(128, activation='relu')(flatten_text)  # Reduced from 256
     
     # Combine features
     combined = Concatenate()([cnn_features, text_features])
     
-    # Fully connected layers
-    x = Dense(512, activation='relu')(combined)
-    x = Dense(256, activation='relu')(x)
+    # Simplified fully connected layers
+    x = Dense(256, activation='relu')(combined)  # Reduced from 512
     output = Dense(num_classes, activation='softmax')(x)
     
     model = Model(inputs=[image_input, text_input], outputs=output)
@@ -122,22 +125,22 @@ def train_model():
     
     # Create and compile model
     model = create_multimodal_model(
-        num_words=len(tokenizer.word_index) + 1,
+        num_words=10000,  # Limited vocabulary size
         num_classes=len(label_encoder.classes_)
     )
     
     model.compile(
-        optimizer='adam',
+        optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
         loss='sparse_categorical_crossentropy',
         metrics=['accuracy']
     )
     
-    # Train model
+    # Train model with reduced epochs
     history = model.fit(
         [image_data, text_data_padded],
         encoded_labels,
         batch_size=BATCH_SIZE,
-        epochs=10,
+        epochs=3,  # Reduced from 10
         validation_split=0.2
     )