Improve model perfomance by stored body coordinates...

Dockerfile

@@ -2,12 +2,19 @@ FROM python:3.10.12-slim
 
 WORKDIR /app
 
-# Set cache directories
+# Set cache directories and performance optimizations
 ENV HF_HOME=/app/.cache
 ENV MPLCONFIGDIR=/app/.cache
+ENV OMP_NUM_THREADS=4
+ENV TOKENIZERS_PARALLELISM=false
 
-# Install dependencies
-RUN apt-get update && apt-get install -y libgl1-mesa-glx libglib2.0-0 && apt-get clean
+# Install dependencies including CUDA support
+RUN apt-get update && apt-get install -y \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    libgomp1 \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
 
 # Create cache directory
 RUN mkdir -p /app/.cache && chmod -R 777 /app/.cache
@@ -27,5 +34,5 @@ COPY . .
 # Set port
 ENV PORT=7860
 
-# Run with Gunicorn
-CMD ["gunicorn", "app:app", "--bind", "0.0.0.0:7860", "--workers", "2"]
+# Run with Gunicorn with optimized settings
+CMD ["gunicorn", "app:app", "--bind", "0.0.0.0:7860", "--workers", "1", "--threads", "4", "--worker-class", "gthread", "--timeout", "300"]

app.py

@@ -1,12 +1,15 @@
-from flask import Flask, render_template, request, send_from_directory
+from flask import Flask, render_template, request, send_from_directory, session
 from PIL import Image
 import os, torch, cv2, mediapipe as mp
 from transformers import SamModel, SamProcessor, logging as hf_logging
 from torchvision import transforms
 from diffusers.utils import load_image
 from flask_cors import CORS 
+import json
+import time
 
 app= Flask(__name__)
+app.secret_key = os.environ.get('SECRET_KEY', 'dev-secret-key-change-in-production')  # Change this to a random secret key
 CORS(app)
 
 # Enable Hugging Face detailed logs (shows model download progress)
@@ -25,16 +28,52 @@ if not os.path.exists(OUTPUT_FOLDER):
     os.makedirs(OUTPUT_FOLDER, exist_ok=True)
 
 
-# Lazy-load model
+# Global model variables
 model, processor = None, None
+device = None
+
+def initialize_model():
+    """Initialize model once at startup"""
+    global model, processor, device
+    
+    # Determine device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"[INFO] Using device: {device}")
+    
+    print("[INFO] Loading SAM model and processor...")
+    model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-50", cache_dir="/app/.cache")
+    processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-50", cache_dir="/app/.cache")
+    
+    # Move model to device
+    model = model.to(device)
+    print(f"[INFO] Model and processor loaded successfully on {device}!")
 
 def load_model():
+    """Ensure model is loaded (should already be loaded at startup)"""
     global model, processor
     if model is None or processor is None:
-        print("[INFO] Loading SAM model and processor...")
-        model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-50", cache_dir="/app/.cache")
-        processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-50", cache_dir="/app/.cache")
-        print("[INFO] Model and processor loaded successfully!")
+        print("[WARNING] Model not loaded, initializing now...")
+        initialize_model()
+
+def warmup_model():
+    """Warm up the model with a dummy inference"""
+    global model, processor, device
+    if model is None or processor is None:
+        return
+    
+    print("[INFO] Warming up model...")
+    try:
+        # Create a dummy image and points for warmup
+        dummy_img = Image.new('RGB', (512, 512), color='white')
+        dummy_points = [[[256, 256], [300, 300]]]
+        inputs = processor(dummy_img, input_points=dummy_points, return_tensors="pt")
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            _ = model(**inputs)
+        print("[INFO] Model warmup completed!")
+    except Exception as e:
+        print(f"[WARNING] Model warmup failed: {e}")
 
 @app.before_request
 def log_request_info():
@@ -49,44 +88,92 @@ def health():
 def serve_output(filename):
     return send_from_directory(OUTPUT_FOLDER, filename)
 
+# Route to serve cached person images
+@app.route('/uploads/<filename>')
+def serve_upload(filename):
+    return send_from_directory(UPLOAD_FOLDER, filename)
+
+def detect_pose_and_get_coordinates(person_path):
+    """Extract pose coordinates from person image"""
+    mp_pose = mp.solutions.pose
+    pose = mp_pose.Pose()
+    image = cv2.imread(person_path)
+    if image is None:
+        raise Exception("No image detected.")
+    
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    results = pose.process(image_rgb)
+    if not results.pose_landmarks:
+        raise Exception("No pose detected.")
+    
+    height, width, _ = image.shape
+    landmarks = results.pose_landmarks.landmark
+    left_shoulder = (int(landmarks[11].x * width), int(landmarks[11].y * height))
+    right_shoulder = (int(landmarks[12].x * width), int(landmarks[12].y * height))
+    
+    return left_shoulder, right_shoulder
+
 @app.route('/', methods=['GET', 'POST'])
 def index():
+    start_time = time.time()
     print(f"[INFO] Handling {request.method} on /")
     if request.method == 'POST':
         try:
             load_model()
-
-            # Save uploaded images
-            person_file = request.files['person_image']
+            
+            # Check if we have a cached person image and coordinates
+            use_cached_person = 'person_coordinates' in session and 'person_image_path' in session
+            person_coordinates = None
+            person_path = None
+            
+            if use_cached_person:
+                # Use cached person image and coordinates
+                person_path = session['person_image_path']
+                person_coordinates = session['person_coordinates']
+                print(f"[INFO] Using cached person image: {person_path}")
+                print(f"[INFO] Using cached coordinates: {person_coordinates}")
+            else:
+                # Process new person image
+                person_file = request.files.get('person_image')
+                if not person_file or person_file.filename == '':
+                    return "No person image provided. Please upload a person image first."
+                
+                person_path = os.path.join(UPLOAD_FOLDER, 'person.jpg')
+                person_file.save(person_path)
+                print(f"[INFO] Saved new person image to {person_path}")
+                
+                # Detect pose and get coordinates
+                left_shoulder, right_shoulder = detect_pose_and_get_coordinates(person_path)
+                person_coordinates = {
+                    'left_shoulder': left_shoulder,
+                    'right_shoulder': right_shoulder
+                }
+                
+                # Cache the person image and coordinates
+                session['person_image_path'] = person_path
+                session['person_coordinates'] = person_coordinates
+                print(f"[INFO] Cached person coordinates: {person_coordinates}")
+
+            # Process garment image
             tshirt_file = request.files['tshirt_image']
-            person_path = os.path.join(UPLOAD_FOLDER, 'person.jpg')
             tshirt_path = os.path.join(UPLOAD_FOLDER, 'tshirt.png')
-            person_file.save(person_path)
             tshirt_file.save(tshirt_path)
-            print(f"[INFO] Saved files to {UPLOAD_FOLDER}")
-
-            # Pose detection
-            mp_pose = mp.solutions.pose
-            pose = mp_pose.Pose()
-            image = cv2.imread(person_path)
-            if image is None:
-                return "No image detected."
-            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            results = pose.process(image_rgb)
-            if not results.pose_landmarks:
-                return "No pose detected."
-            height, width, _ = image.shape
-            landmarks = results.pose_landmarks.landmark
-            left_shoulder = (int(landmarks[11].x * width), int(landmarks[11].y * height))
-            right_shoulder = (int(landmarks[12].x * width), int(landmarks[12].y * height))
-            print(f"[INFO] Shoulder coordinates: {left_shoulder}, {right_shoulder}")
-
-            # SAM model inference
+            print(f"[INFO] Saved garment image to {tshirt_path}")
+
+            # SAM model inference using cached or new coordinates
             img = load_image(person_path)
             new_tshirt = load_image(tshirt_path)
-            input_points = [[[left_shoulder[0], left_shoulder[1]], [right_shoulder[0], right_shoulder[1]]]]
+            input_points = [[[person_coordinates['left_shoulder'][0], person_coordinates['left_shoulder'][1]], 
+                           [person_coordinates['right_shoulder'][0], person_coordinates['right_shoulder'][1]]]]
             inputs = processor(img, input_points=input_points, return_tensors="pt")
-            outputs = model(**inputs)
+            
+            # Move inputs to device
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            # Run inference
+            with torch.no_grad():  # Disable gradient computation for inference
+                outputs = model(**inputs)
+            
             masks = processor.image_processor.post_process_masks(
                 outputs.pred_masks.cpu(),
                 inputs["original_sizes"].cpu(),
@@ -102,8 +189,16 @@ def index():
             img_with_new_tshirt.save(result_path)
             print(f"[INFO] Result saved to {result_path}")
 
-            # Serve via dynamic route
-            return render_template('index.html', result_img='/outputs/result.jpg')
+            # Calculate processing time
+            processing_time = time.time() - start_time
+            print(f"[PERF] Total processing time: {processing_time:.2f}s")
+            
+            # Serve via dynamic route with cached person info
+            return render_template('index.html', 
+                                 result_img='/outputs/result.jpg',
+                                 cached_person=use_cached_person,
+                                 person_image_path=person_path,
+                                 processing_time=f"{processing_time:.2f}s")
 
         except Exception as e:
             print(f"[ERROR] {e}")
@@ -111,7 +206,21 @@ def index():
 
     return render_template('index.html')
 
+@app.route('/change_person', methods=['POST'])
+def change_person():
+    """Clear cached person data to allow new person upload"""
+    session.pop('person_coordinates', None)
+    session.pop('person_image_path', None)
+    print("[INFO] Cleared cached person data")
+    return render_template('index.html')
+
 if __name__ == '__main__':
-        
+    # Initialize model at startup
+    print("[INFO] Initializing model...")
+    initialize_model()
+    
+    # Warm up the model
+    warmup_model()
+    
     print("[INFO] Starting Flask server...")
     app.run(debug=True, host='0.0.0.0')

requirements.txt

@@ -8,4 +8,5 @@ mediapipe
 transformers
 diffusers
 safetensors
-flask-cors
+flask-cors
+accelerate

templates/index.html

@@ -21,7 +21,31 @@
             <div class="grid grid-cols-1 gap-8">
                 <!-- Person Image Upload -->
                 <div>
-                    <h2 class="text-lg font-semibold mb-2">Upload your photo</h2>
+                    <div class="flex justify-between items-center mb-2">
+                        <h2 class="text-lg font-semibold">Upload your photo</h2>
+                        {% if cached_person %}
+                        <button type="button" onclick="changePerson()" class="bg-yellow-500 hover:bg-yellow-600 text-white text-sm px-3 py-1 rounded-lg transition">
+                            Change Person
+                        </button>
+                        {% endif %}
+                    </div>
+                    
+                    {% if cached_person %}
+                    <!-- Show cached person image -->
+                    <div class="border-2 border-green-500 rounded-xl p-4 bg-green-900/20">
+                        <div class="flex flex-col items-center">
+                            <div class="flex items-center gap-2 mb-2">
+                                <svg xmlns="http://www.w3.org/2000/svg" class="h-5 w-5 text-green-400" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+                                    <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M5 13l4 4L19 7" />
+                                </svg>
+                                <p class="text-green-400 text-sm font-medium">Person image cached</p>
+                            </div>
+                            <img src="/uploads/person.jpg" alt="Cached Person" class="max-h-32 rounded-lg border border-gray-600">
+                            <p class="text-gray-400 text-xs mt-2">Coordinates saved - no need to re-upload!</p>
+                        </div>
+                    </div>
+                    {% else %}
+                    <!-- Upload new person image -->
                     <label for="person_image" class="flex flex-col items-center justify-center border-2 border-dashed border-gray-600 rounded-xl p-6 hover:bg-gray-700 cursor-pointer">
                         <svg xmlns="http://www.w3.org/2000/svg" class="h-10 w-10 text-gray-400 mb-2" fill="none"
                             viewBox="0 0 24 24" stroke="currentColor">
@@ -29,18 +53,29 @@
                                 d="M7 16v4m0 0h10m-10 0v-4m0 0h10m-10 0V5m0 0h10m-10 0H5m14 0h-2" />
                         </svg>
                         <p class="text-gray-400">Drag & drop or click to upload</p>
-                        <input id="person_image" type="file" name="person_image" class="hidden" required
+                        <input id="person_image" type="file" name="person_image" class="hidden" 
                             onchange="showFileName('person_image', 'person_filename', 'person_preview')">
                     </label>
                     <p id="person_filename" class="text-green-400 text-sm mt-2 text-center"></p>
                     <div class="mt-3 flex justify-center">
                         <img id="person_preview" class="hidden max-h-32 rounded-lg border border-gray-600">
                     </div>
+                    {% endif %}
                 </div>
 
                 <!-- Garment Image Upload with Cropper -->
                 <div>
-                    <h2 class="text-lg font-semibold mb-2">Upload garment image</h2>
+                    <div class="flex justify-between items-center mb-2">
+                        <h2 class="text-lg font-semibold">Upload garment image</h2>
+                        {% if cached_person %}
+                        <div class="flex items-center gap-1 text-xs text-green-400">
+                            <svg xmlns="http://www.w3.org/2000/svg" class="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+                                <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M13 10V3L4 14h7v7l9-11h-7z" />
+                            </svg>
+                            <span>Fast mode</span>
+                        </div>
+                        {% endif %}
+                    </div>
                     <label for="tshirt_image" class="flex flex-col items-center justify-center border-2 border-dashed border-gray-600 rounded-xl p-6 hover:bg-gray-700 cursor-pointer">
                         <svg xmlns="http://www.w3.org/2000/svg" class="h-10 w-10 text-gray-400 mb-2" fill="none"
                             viewBox="0 0 24 24" stroke="currentColor">
@@ -73,6 +108,13 @@
             {% if result_img %}
             <div>
                 <h2 class="text-2xl font-bold mb-6 text-center">🎉 Your Virtual Try-On Result</h2>
+                {% if processing_time %}
+                <div class="text-center mb-4">
+                    <span class="bg-blue-500 text-white px-3 py-1 rounded-full text-sm">
+                        ⚡ Processed in {{ processing_time }}
+                    </span>
+                </div>
+                {% endif %}
                 <div class="flex justify-center mb-6">
                     <img id="result-image" src="{{ result_img }}" alt="Result Image" class="rounded-xl">
                 </div>
@@ -139,6 +181,25 @@
             document.getElementById('loading-spinner').style.display = 'flex';
         });
 
+        // Change Person function
+        function changePerson() {
+            fetch('/change_person', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/x-www-form-urlencoded',
+                }
+            })
+            .then(response => response.text())
+            .then(html => {
+                // Replace the current page content
+                document.documentElement.innerHTML = html;
+            })
+            .catch(error => {
+                console.error('Error changing person:', error);
+                alert('Error changing person. Please refresh the page.');
+            });
+        }
+
         // Show file name + preview (person only)
         function showFileName(inputId, filenameId, previewId) {
             const input = document.getElementById(inputId);