app.py

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw

print("🚀 Starting AI Headshot Generator...")

# Check for GPU WITHOUT crashing
try:
    import torch
    HAS_TORCH = True
    print(f"✅ PyTorch {torch.__version__} available")
    
    if hasattr(torch, 'cuda') and torch.cuda.is_available():
        HAS_CUDA = True
        print(f"✅ CUDA available")
        print(f"✅ GPU: {torch.cuda.get_device_name(0)}")
    else:
        HAS_CUDA = False
        print("⚠️ CUDA not available - running on CPU")
except:
    HAS_TORCH = False
    HAS_CUDA = False
    print("⚠️ PyTorch not available")

# ===== LICENSE SYSTEM =====
ADMIN_KEYS = {"HEADSHOT-TEST123456", "HEADSHOT-OWNERACCESS"}
MAX_TRIALS = 3

def verify_license(key):
    return key.strip().upper() in ADMIN_KEYS if key else False

# ===== SIMPLE GENERATOR =====
def generate_headshot(image, license_key):
    """Generate placeholder headshot"""
    
    # Verify license
    if not verify_license(license_key):
        return np.zeros((512, 512, 3), dtype=np.uint8), "❌ Invalid license. Use: HEADSHOT-TEST123456"
    
    # Create result image
    if image is not None:
        h, w = image.shape[:2]
        result = np.ones((h, w, 3), dtype=np.uint8) * 240
    else:
        result = np.ones((512, 512, 3), dtype=np.uint8) * 240
    
    # Add info
    img = Image.fromarray(result)
    draw = ImageDraw.Draw(img)
    
    draw.text((30, 30), "🎯 AI Headshot Generator", fill=(0, 0, 0), size=20)
    draw.text((30, 70), f"Hardware: {'GPU' if HAS_CUDA else 'CPU'}", fill=(50, 50, 150))
    draw.text((30, 100), "Test Key: HEADSHOT-TEST123456", fill=(0, 100, 0))
    draw.text((30, 130), "Status: Ready to generate", fill=(0, 150, 0))
    draw.text((30, 160), "Next: Add AI models", fill=(150, 50, 0))
    
    return np.array(img), f"✅ Generated on {'GPU' if HAS_CUDA else 'CPU'}"

# ===== GRADIO UI =====
with gr.Blocks(title="AI Headshot Generator", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown(f"""
    # 🎯 AI Headshot Generator
    *Running on: {'🚀 GPU' if HAS_CUDA else '⚡ CPU'}*
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Hardware Status
            gr.Markdown(f"""
            ### ⚙️ Hardware Status
            - **PyTorch**: {'✅ Available' if HAS_TORCH else '❌ Not available'}
            - **CUDA**: {'✅ Available' if HAS_CUDA else '⚠️ CPU mode'}
            """)
            
            # Upload
            image_input = gr.Image(
                label="📸 Upload Photo",
                type="numpy",
                height=250
            )
            
            # License
            license_input = gr.Textbox(
                label="🔑 License Key",
                value="HEADSHOT-TEST123456",
                placeholder="Enter HEADSHOT-TEST123456"
            )
            
            # Generate Button
            generate_btn = gr.Button(
                "✨ Generate Headshot",
                variant="primary",
                size="lg"
            )
            
            # Status
            status_output = gr.Textbox(
                label="Status",
                interactive=False
            )
            
            gr.Markdown(f"""
            ### 🎯 Free Trials
            - **Test Key**: `HEADSHOT-TEST123456`
            - **Trials**: {MAX_TRIALS} free generations
            - **Purchase**: [Get License](https://canadianheadshotpro.carrd.co)
            """)
        
        with gr.Column(scale=1):
            # Result
            result_image = gr.Image(
                label="🖼️ Your Headshot",
                height=400
            )
            
            gr.Markdown("""
            ### 📥 Download Options
            - **Free Trial**: Watermarked preview
            - **Premium**: HD quality, no watermark
            - **Professional**: Multiple styles
            
            ⏱️ *Generation time: 10-20 seconds*
            """)

    # Connect
    generate_btn.click(
        fn=generate_headshot,
        inputs=[image_input, license_input],
        outputs=[result_image, status_output]
    )

if __name__ == "__main__":
    # Remove share=True for Hugging Face Spaces
    demo.launch()