Update app.py

app.py

@@ -6,11 +6,13 @@ import os
 import time
 import secrets
 import json
+import hashlib
+from datetime import datetime, timedelta
 
 import spaces
 
 import PIL
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from typing import Tuple
 
 import diffusers
@@ -91,6 +93,109 @@ def verify_license(license_key):
     valid_licenses = load_licenses()
     return license_upper in valid_licenses
 
+# ===== TRIAL SYSTEM =====
+TRIAL_FILE = "user_trials.json"
+MAX_FREE_TRIALS = 3
+
+def load_trials():
+    """Load user trial data"""
+    try:
+        with open(TRIAL_FILE, 'r') as f:
+            return json.load(f)
+    except FileNotFoundError:
+        return {}
+
+def save_trials(trials_data):
+    """Save user trial data"""
+    with open(TRIAL_FILE, 'w') as f:
+        json.dump(trials_data, f)
+
+def get_user_identifier():
+    """Create a unique but anonymous user identifier"""
+    # For Hugging Face Spaces, we'll use a simple approach
+    # In production, you might want to use session-based tracking
+    return "gradio_user"
+
+def can_use_free_trial(user_id):
+    """Check if user can use free trial"""
+    trials_data = load_trials()
+    
+    if user_id not in trials_data:
+        return True, MAX_FREE_TRIALS
+    
+    user_data = trials_data[user_id]
+    trials_used = user_data.get('trials_used', 0)
+    first_trial_date = user_data.get('first_trial_date')
+    
+    # Reset trials after 30 days
+    if first_trial_date:
+        first_date = datetime.fromisoformat(first_trial_date)
+        if datetime.now() - first_date > timedelta(days=30):
+            trials_used = 0
+            user_data['trials_used'] = 0
+            user_data['first_trial_date'] = datetime.now().isoformat()
+            save_trials(trials_data)
+    
+    trials_left = MAX_FREE_TRIALS - trials_used
+    return trials_left > 0, trials_left
+
+def record_trial_usage(user_id):
+    """Record that user used a trial"""
+    trials_data = load_trials()
+    
+    if user_id not in trials_data:
+        trials_data[user_id] = {
+            'trials_used': 1,
+            'first_trial_date': datetime.now().isoformat(),
+            'last_used': datetime.now().isoformat()
+        }
+    else:
+        trials_data[user_id]['trials_used'] += 1
+        trials_data[user_id]['last_used'] = datetime.now().isoformat()
+    
+    save_trials(trials_data)
+
+def apply_watermark(image):
+    """Apply watermark to free trial images"""
+    # Convert to PIL if needed
+    if hasattr(image, 'mode'):
+        pil_image = image
+    else:
+        pil_image = Image.fromarray(image)
+    
+    draw = ImageDraw.Draw(pil_image, 'RGBA')
+    width, height = pil_image.size
+    
+    # Watermark text
+    watermark_text = "PREVIEW - UPGRADE TO DOWNLOAD"
+    
+    # Use default font
+    try:
+        font = ImageFont.truetype("arial.ttf", min(width, height) // 20)
+    except:
+        try:
+            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", min(width, height) // 20)
+        except:
+            font = ImageFont.load_default()
+    
+    # Get text size
+    bbox = draw.textbbox((0, 0), watermark_text, font=font)
+    text_width = bbox[2] - bbox[0]
+    text_height = bbox[3] - bbox[1]
+    
+    # Position watermark (center bottom)
+    x = (width - text_width) // 2
+    y = height - text_height - 50
+    
+    # Draw semi-transparent background
+    draw.rectangle([x-10, y-10, x+text_width+10, y+text_height+10], 
+                   fill=(0, 0, 0, 128))
+    
+    # Draw text
+    draw.text((x, y), watermark_text, fill=(255, 255, 255, 255), font=font)
+    
+    return pil_image
+
 # Initialize licenses on first run
 VALID_LICENSES = load_licenses()
 print(f"✅ License system initialized. Admin keys: {ADMIN_KEYS}")
@@ -278,7 +383,7 @@ def save_as_png(image, filename="professional_headshot"):
 @spaces.GPU
 def generate_image(
     face_image_path,
-    license_key,  # NEW: License key parameter
+    license_key,
     prompt,
     negative_prompt,
     style_name,
@@ -296,15 +401,27 @@ def generate_image(
     progress=gr.Progress(track_tqdm=True),
 ):
 
-    # ===== LICENSE VERIFICATION =====
-    if not verify_license(license_key):
-        raise gr.Error("""❌ Invalid or missing license key. 
-
-Please check your email for the license key sent after purchase. 
-If you haven't purchased yet, visit: 
-https://canadianheadshotpro.carrd.co
-
-Need help? Email: bee.tools@zohomailcloud.ca""")
+    # ===== LICENSE & TRIAL VERIFICATION =====
+    user_id = get_user_identifier()
+    has_valid_license = verify_license(license_key)
+    
+    # If no valid license, check free trials
+    if not has_valid_license:
+        can_use_trial, trials_left = can_use_free_trial(user_id)
+        
+        if not can_use_trial:
+            raise gr.Error(f"""
+            ❌ No free trials remaining! 
+            
+            You've used all {MAX_FREE_TRIALS} free generations. 
+            
+            🔑 **Options:**
+            1. **Purchase a license** for unlimited HD downloads
+            2. **Enter your existing license key** if you already purchased
+            
+            💡 Visit: https://canadianheadshotpro.carrd.co to purchase
+            📧 Support: bee.tools@zohomailcloud.ca
+            """)
 
     if enable_LCM:
         pipe.scheduler = diffusers.LCMScheduler.from_config(pipe.scheduler.config)
@@ -393,7 +510,21 @@ Need help? Email: bee.tools@zohomailcloud.ca""")
         generator=generator,
     ).images
 
-    png_filepath = save_as_png(images[0])
+    # ===== APPLY WATERMARK IF USING FREE TRIAL =====
+    final_image = images[0]
+    if not has_valid_license:
+        # Record trial usage
+        record_trial_usage(user_id)
+        # Apply watermark
+        final_image = apply_watermark(final_image)
+        
+        # Update trials left
+        _, trials_left = can_use_free_trial(user_id)
+        
+        # Show trial usage message
+        gr.Info(f"Free trial used! {trials_left} generations remaining. Upgrade for watermark-free HD downloads.")
+
+    png_filepath = save_as_png(final_image)
     return png_filepath, gr.update(visible=True)
 
 # MODERN PROFESSIONAL UI CSS
@@ -413,13 +544,11 @@ css = """
     --border: #e2e8f0;
     --shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
 }
-
 .gradio-container {
     max-width: 1400px !important;
     font-family: 'Inter', 'SF Pro Display', -apple-system, BlinkMacSystemFont, sans-serif;
     background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%) !important;
 }
-
 .main-container {
     background: white;
     border-radius: 24px;
@@ -428,7 +557,6 @@ css = """
     overflow: hidden;
     border: 1px solid var(--border);
 }
-
 .hero-section {
     background: linear-gradient(135deg, var(--primary) 0%, var(--secondary) 100%);
     color: white;
@@ -437,7 +565,6 @@ css = """
     position: relative;
     overflow: hidden;
 }
-
 .hero-section::before {
     content: '';
     position: absolute;
@@ -448,7 +575,6 @@ css = """
     background: url('data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 100" fill="rgba(255,255,255,0.1)"><polygon points="0,0 1000,100 0,100"/></svg>');
     background-size: cover;
 }
-
 .hero-title {
     font-size: 3em;
     font-weight: 800;
@@ -458,7 +584,6 @@ css = """
     -webkit-text-fill-color: transparent;
     background-clip: text;
 }
-
 .hero-subtitle {
     font-size: 1.3em;
     font-weight: 400;
@@ -466,7 +591,6 @@ css = """
     max-width: 600px;
     margin: 0 auto;
 }
-
 .upload-area {
     border: 3px dashed var(--border);
     border-radius: 20px;
@@ -476,19 +600,16 @@ css = """
     transition: all 0.3s ease;
     cursor: pointer;
 }
-
 .upload-area:hover {
     border-color: var(--primary);
     background: #f0f9ff;
     transform: translateY(-2px);
 }
-
 .upload-icon {
     font-size: 3em;
     margin-bottom: 16px;
     color: var(--primary);
 }
-
 .control-card {
     background: var(--surface);
     border-radius: 16px;
@@ -498,18 +619,15 @@ css = """
     box-shadow: 0 2px 8px rgba(0, 0, 0, 0.04);
     transition: all 0.3s ease;
 }
-
 .control-card:hover {
     box-shadow: 0 4px 12px rgba(0, 0, 0, 0.08);
     transform: translateY(-1px);
 }
-
 .control-header {
     display: flex;
     align-items: center;
     margin-bottom: 16px;
 }
-
 .control-icon {
     width: 40px;
     height: 40px;
@@ -522,14 +640,12 @@ css = """
     color: white;
     font-weight: 600;
 }
-
 .control-title {
     font-size: 1.2em;
     font-weight: 600;
     color: var(--text-primary);
     margin: 0;
 }
-
 .result-card {
     background: var(--surface);
     border-radius: 20px;
@@ -538,24 +654,20 @@ css = """
     box-shadow: var(--shadow);
     height: 100%;
 }
-
 .result-header {
     text-align: center;
     margin-bottom: 24px;
 }
-
 .result-title {
     font-size: 1.5em;
     font-weight: 700;
     color: var(--text-primary);
     margin-bottom: 8px;
 }
-
 .result-subtitle {
     color: var(--text-secondary);
     font-size: 0.95em;
 }
-
 .image-container {
     border-radius: 16px;
     overflow: hidden;
@@ -563,7 +675,6 @@ css = """
     border: 1px solid var(--border);
     margin-bottom: 20px;
 }
-
 .success-banner {
     background: linear-gradient(135deg, var(--success), #059669);
     color: white;
@@ -572,7 +683,6 @@ css = """
     margin-top: 20px;
     text-align: center;
 }
-
 .btn-primary {
     background: linear-gradient(135deg, var(--primary), var(--primary-dark)) !important;
     color: white !important;
@@ -584,23 +694,19 @@ css = """
     transition: all 0.3s ease !important;
     box-shadow: 0 4px 12px rgba(37, 99, 235, 0.3) !important;
 }
-
 .btn-primary:hover {
     transform: translateY(-2px) !important;
     box-shadow: 0 6px 20px rgba(37, 99, 235, 0.4) !important;
 }
-
 .slider-container {
     padding: 8px 0;
 }
-
 .slider-label {
     display: flex;
     justify-content: space-between;
     align-items: center;
     margin-bottom: 8px;
 }
-
 .slider-value {
     background: var(--primary);
     color: white;
@@ -609,7 +715,6 @@ css = """
     font-size: 0.85em;
     font-weight: 600;
 }
-
 .tips-card {
     background: linear-gradient(135deg, #fef3c7, #f59e0b);
     border: none;
@@ -617,29 +722,48 @@ css = """
     padding: 20px;
     margin-bottom: 20px;
 }
-
 .tips-header {
     display: flex;
     align-items: center;
     margin-bottom: 12px;
 }
-
 .tips-icon {
     font-size: 1.5em;
     margin-right: 12px;
 }
-
 .progress-container {
     margin: 20px 0;
     text-align: center;
 }
-
 .progress-text {
     font-size: 0.9em;
     color: var(--text-secondary);
     margin-top: 8px;
 }
-
+.trial-banner {
+    background: linear-gradient(135deg, #10b981, #059669);
+    color: white;
+    padding: 15px;
+    border-radius: 12px;
+    text-align: center;
+    margin-bottom: 15px;
+}
+.trial-banner-warning {
+    background: linear-gradient(135deg, #f59e0b, #d97706);
+    color: white;
+    padding: 15px;
+    border-radius: 12px;
+    text-align: center;
+    margin-bottom: 15px;
+}
+.trial-banner-error {
+    background: linear-gradient(135deg, #ef4444, #dc2626);
+    color: white;
+    padding: 15px;
+    border-radius: 12px;
+    text-align: center;
+    margin-bottom: 15px;
+}
 /* Responsive design */
 @media (max-width: 768px) {
     .hero-title {
@@ -697,24 +821,36 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                         elem_classes="upload-area"
                     )
                 
-                # License Key Section - NEW
+                # Access Options Section - UPDATED WITH TRIAL SYSTEM
                 with gr.Column(elem_classes="control-card"):
                     gr.HTML("""
                     <div class="control-header">
                         <div class="control-icon">🔑</div>
-                        <h3 class="control-title">License Verification</h3>
+                        <h3 class="control-title">Access Options</h3>
+                    </div>
+                    """)
+                    
+                    # Trial Status Display
+                    trial_status = gr.HTML(f"""
+                    <div class="trial-banner">
+                        <h4 style="margin: 0 0 8px 0;">🎉 {MAX_FREE_TRIALS} FREE Trials Available!</h4>
+                        <p style="margin: 0; font-size: 0.9em;">Try our AI headshot generator - no credit card required</p>
                     </div>
                     """)
+                    
                     license_input = gr.Textbox(
                         label="",
-                        placeholder="Enter your license key (e.g., HEADSHOT-A1B2C3D4E5F6)",
+                        placeholder="Enter license key (or leave blank for free trial)",
                         show_label=False,
-                        info="💡 Your license key was emailed to you after purchase"
+                        info=f"💡 You get {MAX_FREE_TRIALS} free generations. Purchase license for HD downloads without watermark."
                     )
-                    gr.HTML("""
+                    
+                    gr.HTML(f"""
                     <div style="font-size: 0.85em; color: var(--text-secondary); margin-top: 8px;">
-                        <strong>Test Keys:</strong> HEADSHOT-TEST123456, HEADSHOT-OWNERACCESS
-                        <br>Don't have a license? <a href="https://canadianheadshotpro.carrd.co" target="_blank">Purchase access here</a>
+                        <strong>Free Trial:</strong> {MAX_FREE_TRIALS} watermarked previews<br>
+                        <strong>Premium License:</strong> Unlimited HD downloads, no watermark<br>
+                        <strong>Professional Use:</strong> Commercial rights included<br>
+                        <a href="https://canadianheadshotpro.carrd.co" target="_blank" style="color: var(--primary); font-weight: 600;">👉 Click here to purchase license</a>
                     </div>
                     """)
                 
@@ -917,11 +1053,39 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
         </div>
         """
 
+    def update_trial_display(license_key):
+        """Update trial counter based on license status"""
+        if verify_license(license_key):
+            return """
+            <div class="trial-banner">
+                <h4 style="margin: 0 0 8px 0;">✅ Premium License Active</h4>
+                <p style="margin: 0; font-size: 0.9em;">Unlimited HD downloads - no watermark</p>
+            </div>
+            """
+        
+        user_id = get_user_identifier()
+        can_use, trials_left = can_use_free_trial(user_id)
+        
+        if not can_use:
+            return f"""
+            <div class="trial-banner-error">
+                <h4 style="margin: 0 0 8px 0;">❌ No Free Trials Left</h4>
+                <p style="margin: 0; font-size: 0.9em;">Please purchase a license to continue</p>
+            </div>
+            """
+        
+        return f"""
+        <div class="trial-banner">
+            <h4 style="margin: 0 0 8px 0;">🎉 {trials_left} Free Generations Left!</h4>
+            <p style="margin: 0; font-size: 0.9em;">Watermarked previews - upgrade for HD downloads</p>
+        </div>
+        """
+
     submit.click(
         fn=generate_image,
         inputs=[
             face_file,
-            license_input,  # NEW: Added license input
+            license_input,
             prompt,
             negative_prompt,
             style,
@@ -943,6 +1107,13 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
         outputs=success_msg
     )
 
+    # Update trial display when license input changes
+    license_input.change(
+        fn=update_trial_display,
+        inputs=[license_input],
+        outputs=[trial_status]
+    )
+
     enable_LCM.input(
         fn=toggle_lcm_ui,
         inputs=[enable_LCM],