Update app.py

app.py

@@ -1,1268 +1,123 @@
 import os
-os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
-os.environ['CUDA_MODULE_LOADING'] = 'LAZY'
-
-import cv2
-import torch
-import random
-import numpy as np
-import time
-import secrets
-import json
-import hashlib
-from datetime import datetime, timedelta
-
-#import spaces
-
-import PIL
-from PIL import Image, ImageDraw, ImageFont
-from typing import Tuple
-
-import diffusers
-from diffusers.utils import load_image
-from diffusers.models import ControlNetModel
-from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
-
-from huggingface_hub import hf_hub_download
-
-from insightface.app import FaceAnalysis
-
-# Optional imports with fallbacks
-try:
-    from style_template import styles
-    STYLE_NAMES = list(styles.keys())
-except:
-    styles = {}
-    STYLE_NAMES = []
-
-try:
-    from pipeline_stable_diffusion_xl_instantid_full import StableDiffusionXLInstantIDPipeline, draw_kps
-except:
-    # Create dummy functions if imports fail
-    class StableDiffusionXLInstantIDPipeline:
-        def __init__(self, *args, **kwargs):
-            pass
-    def draw_kps(*args, **kwargs):
-        return None
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 
 import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw
+import torch
 
-try:
-    from depth_anything.dpt import DepthAnything
-    from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
-    import torch.nn.functional as F
-    from torchvision.transforms import Compose
-    DEPTH_AVAILABLE = True
-except:
-    DEPTH_AVAILABLE = False
-
-# GPU optimization
-torch.backends.cudnn.benchmark = True
-torch.backends.cuda.matmul.allow_tf32 = True
-
-# global variable
-MAX_SEED = np.iinfo(np.int32).max
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32
-enable_lcm_arg = False
-
-print(f"🚀 Starting AI Headshot Generator on {device}...")
+print(f"🚀 Starting on ZeroGPU with PyTorch {torch.__version__}")
+print(f"✅ CUDA available: {torch.cuda.is_available()}")
+print(f"✅ GPU: {torch.cuda.get_device_name(0)}")
 
 # ===== LICENSE SYSTEM =====
-LICENSE_FILE = "valid_licenses.json"
-
-# PERMANENT ADMIN/TEST KEYS (Always work)
-ADMIN_KEYS = {
-    "HEADSHOT-TEST123456",
-    "HEADSHOT-OWNERACCESS", 
-    "HEADSHOT-DEVELOPER123",
-    "HEADSHOT-ADMIN789012"
-}
-
-def load_licenses():
-    """Load valid licenses from file"""
-    try:
-        with open(LICENSE_FILE, 'r') as f:
-            return set(json.load(f))
-    except FileNotFoundError:
-        # Initialize with admin keys on first run
-        initial_licenses = ADMIN_KEYS.copy()
-        save_licenses(initial_licenses)
-        return initial_licenses
-
-def save_licenses(licenses):
-    """Save licenses to file"""
-    with open(LICENSE_FILE, 'w') as f:
-        json.dump(list(licenses), f)
-
-def generate_license_key():
-    """Generate a unique license key"""
-    license_key = f"HEADSHOT-{secrets.token_hex(6).upper()}"
-    valid_licenses = load_licenses()
-    valid_licenses.add(license_key)
-    save_licenses(valid_licenses)
-    return license_key
-
-def verify_license(license_key):
-    """Check if license key is valid"""
-    if not license_key or not license_key.strip():
-        return False
-    
-    license_upper = license_key.strip().upper()
-    
-    # Check admin keys first (always work)
-    if license_upper in ADMIN_KEYS:
-        print(f"✅ Admin access granted with: {license_upper}")
-        return True
-    
-    # Check purchased licenses
-    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"""
-    # Simple identifier for Spaces
-    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
-VALID_LICENSES = load_licenses()
-print(f"✅ License system initialized. Admin keys available.")
-
-# ===== LAZY LOAD MODELS =====
-print("🔄 Loading AI models (this may take a few minutes)...")
-
-# Create temp directory if needed
-os.makedirs("./checkpoints", exist_ok=True)
-os.makedirs("temp_downloads", exist_ok=True)
-
-# Load models with error handling
-try:
-    # Download checkpoints with timeout
-    print("📥 Downloading InstantID checkpoints...")
-    hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/config.json", local_dir="./checkpoints")
-    hf_hub_download(
-        repo_id="InstantX/InstantID",
-        filename="ControlNetModel/diffusion_pytorch_model.safetensors",
-        local_dir="./checkpoints",
-    )
-    hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="./checkpoints")
-    print("✅ Checkpoints downloaded")
-except Exception as e:
-    print(f"⚠️ Could not download checkpoints: {e}")
-    print("⚠️ App will use fallback mode")
-
-# Load face encoder
-try:
-    app = FaceAnalysis(
-        name="antelopev2",
-        root="./",
-        providers=["CPUExecutionProvider"],
-    )
-    app.prepare(ctx_id=0, det_size=(640, 640))
-    print("✅ Face encoder loaded")
-except Exception as e:
-    print(f"⚠️ Could not load face encoder: {e}")
-    app = None
+ADMIN_KEYS = {"HEADSHOT-TEST123456", "HEADSHOT-OWNERACCESS"}
+MAX_TRIALS = 3
 
-# Initialize depth model if available
-if DEPTH_AVAILABLE:
-    try:
-        depth_anything = DepthAnything.from_pretrained('LiheYoung/depth_anything_vitl14').to(device).eval()
-        transform = Compose([
-            Resize(
-                width=518,
-                height=518,
-                resize_target=False,
-                keep_aspect_ratio=True,
-                ensure_multiple_of=14,
-                resize_method='lower_bound',
-                image_interpolation_method=cv2.INTER_CUBIC,
-            ),
-            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-            PrepareForNet(),
-        ])
-    except:
-        DEPTH_AVAILABLE = False
-        print("⚠️ Depth model not available")
+def verify_license(key):
+    return key.strip().upper() in ADMIN_KEYS if key else False
 
-# Path to InstantID models
-face_adapter = f"./checkpoints/ip-adapter.bin"
-controlnet_path = f"./checkpoints/ControlNetModel"
-
-# Initialize pipeline with error handling
-pipe = None
-controlnet_identitynet = None
-controlnet_canny = None
-controlnet_depth = None
-
-try:
-    print("🔄 Loading ControlNet models...")
-    controlnet_identitynet = ControlNetModel.from_pretrained(
-        controlnet_path, torch_dtype=dtype
-    )
-    
-    controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
-    controlnet_depth_model = "diffusers/controlnet-depth-sdxl-1.0-small"
-    
-    controlnet_canny = ControlNetModel.from_pretrained(
-        controlnet_canny_model, torch_dtype=dtype
-    ).to(device)
+# ===== SIMPLE GENERATOR =====
+def generate_headshot(image, license_key):
+    """Placeholder that shows GPU is working"""
     
-    if DEPTH_AVAILABLE:
-        controlnet_depth = ControlNetModel.from_pretrained(
-            controlnet_depth_model, torch_dtype=dtype
-        ).to(device)
-    
-    print("🔄 Loading Stable Diffusion pipeline...")
-    pretrained_model_name_or_path = "wangqixun/YamerMIX_v8"
-    
-    pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
-        pretrained_model_name_or_path,
-        controlnet=[controlnet_identitynet],
-        torch_dtype=dtype,
-        safety_checker=None,
-        feature_extractor=None,
-    ).to(device)
-    
-    pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(
-        pipe.scheduler.config
-    )
+    # Verify license
+    if not verify_license(license_key):
+        return np.zeros((512, 512, 3), dtype=np.uint8), "❌ Invalid license"
     
-    # Load and disable LCM
-    pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
-    pipe.disable_lora()
-    
-    pipe.cuda()
-    pipe.load_ip_adapter_instantid(face_adapter)
-    
-    if hasattr(pipe, 'image_proj_model'):
-        pipe.image_proj_model.to("cuda")
-    if hasattr(pipe, 'unet'):
-        pipe.unet.to("cuda")
-    
-    print("✅ AI pipeline loaded successfully!")
-    
-except Exception as e:
-    print(f"⚠️ Could not load AI pipeline: {e}")
-    print("⚠️ App will run in demo mode")
-
-# ControlNet functions with fallbacks
-if DEPTH_AVAILABLE:
-    def get_depth_map(image):
-        image = np.array(image) / 255.0
+    # Create result image
+    if image is not None:
         h, w = image.shape[:2]
-        image = transform({'image': image})['image']
-        image = torch.from_numpy(image).unsqueeze(0).to("cuda")
-        with torch.no_grad():
-            depth = depth_anything(image)
-        depth = F.interpolate(depth[None], (h, w), mode='bilinear', align_corners=False)[0, 0]
-        depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
-        depth = depth.cpu().numpy().astype(np.uint8)
-        depth_image = Image.fromarray(depth)
-        return depth_image
-else:
-    def get_depth_map(image):
-        return Image.fromarray(np.zeros((512, 512), dtype=np.uint8))
-
-def get_canny_image(image, t1=100, t2=200):
-    try:
-        image_np = np.array(image)
-        if len(image_np.shape) == 3:
-            image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
-            edges = cv2.Canny(image_np, t1, t2)
-            return Image.fromarray(edges, "L")
-        return Image.fromarray(np.zeros((512, 512), dtype=np.uint8))
-    except:
-        return Image.fromarray(np.zeros((512, 512), dtype=np.uint8))
-
-controlnet_map = {
-    "canny": controlnet_canny,
-    "depth": controlnet_depth,
-}
-controlnet_map_fn = {
-    "canny": get_canny_image,
-    "depth": get_depth_map,
-}
-
-def toggle_lcm_ui(value):
-    if value:
-        return (
-            gr.update(minimum=0, maximum=100, step=1, value=5),
-            gr.update(minimum=0.1, maximum=20.0, step=0.1, value=1.5),
-        )
+        result = np.ones((h, w, 3), dtype=np.uint8) * 240
     else:
-        return (
-            gr.update(minimum=5, maximum=100, step=1, value=30),
-            gr.update(minimum=0.1, maximum=20.0, step=0.1, value=5),
-        )
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
-def convert_from_cv2_to_image(img: np.ndarray) -> Image:
-    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-
-def convert_from_image_to_cv2(img: Image) -> np.ndarray:
-    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-
-def resize_img(
-    input_image,
-    max_side=1280,
-    min_side=1024,
-    size=None,
-    pad_to_max_side=False,
-    mode=PIL.Image.BILINEAR,
-    base_pixel_number=64,
-):
-    w, h = input_image.size
-    if size is not None:
-        w_resize_new, h_resize_new = size
-    else:
-        ratio = min_side / min(h, w)
-        w, h = round(ratio * w), round(ratio * h)
-        ratio = max_side / max(h, w)
-        input_image = input_image.resize([round(ratio * w), round(ratio * h)], mode)
-        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
-        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
-    input_image = input_image.resize([w_resize_new, h_resize_new], mode)
-
-    if pad_to_max_side:
-        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
-        offset_x = (max_side - w_resize_new) // 2
-        offset_y = (max_side - h_resize_new) // 2
-        res[
-            offset_y : offset_y + h_resize_new, offset_x : offset_x + w_resize_new
-        ] = np.array(input_image)
-        input_image = Image.fromarray(res)
-    return input_image
-
-def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    if style_name == "No Style" or style_name not in styles:
-        return positive, negative
-    p, n = styles.get(style_name, ("{prompt}", ""))
-    return p.replace("{prompt}", positive), n + " " + negative
-
-def save_as_png(image, filename="professional_headshot"):
-    """Save image as PNG and return file path for download"""
-    temp_dir = "temp_downloads"
-    os.makedirs(temp_dir, exist_ok=True)
+        result = np.ones((512, 512, 3), dtype=np.uint8) * 240
     
-    timestamp = int(time.time())
-    filepath = os.path.join(temp_dir, f"{filename}_{timestamp}.png")
+    # Add GPU status
+    img = Image.fromarray(result)
+    draw = ImageDraw.Draw(img)
     
-    if image.mode in ('RGBA', 'LA'):
-        background = Image.new('RGB', image.size, (255, 255, 255))
-        background.paste(image, mask=image.split()[-1])
-        image = background
-    elif image.mode != 'RGB':
-        image = image.convert('RGB')
+    draw.text((30, 30), "✅ ZEROGPU WORKING", fill=(0, 150, 0), size=20)
+    draw.text((30, 70), f"PyTorch {torch.__version__}", fill=(0, 0, 0))
+    draw.text((30, 100), f"GPU: {torch.cuda.get_device_name(0)}", fill=(0, 0, 0))
+    draw.text((30, 130), f"Memory: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB", fill=(0, 0, 0))
+    draw.text((30, 160), "Test Key: HEADSHOT-TEST123456", fill=(50, 50, 150))
+    draw.text((30, 190), "Next: Add AI models", fill=(150, 50, 0))
     
-    image.save(filepath, "PNG", optimize=True)
-    return filepath
+    return np.array(img), "✅ Success! GPU is working."
 
-@spaces.GPU
-def generate_image(
-    face_image_path,
-    license_key,
-    prompt,
-    negative_prompt,
-    style_name,
-    num_steps,
-    identitynet_strength_ratio,
-    adapter_strength_ratio,
-    canny_strength,
-    depth_strength,
-    controlnet_selection,
-    guidance_scale,
-    seed,
-    scheduler,
-    enable_LCM,
-    enhance_face_region,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """Generate AI headshot with license and trial system"""
-    
-    # ===== 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. 
+# ===== GRADIO UI =====
+with gr.Blocks(title="AI Headshot Generator", theme=gr.themes.Soft()) as demo:
+    
+    gr.Markdown("""
+    # 🎯 AI Headshot Generator
+    *Powered by ZeroGPU - PyTorch ready*
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # GPU Status
+            gr.Markdown(f"""
+            ### 🚀 ZeroGPU Status
+            - **PyTorch**: {torch.__version__}
+            - **CUDA**: {'✅ Available' if torch.cuda.is_available() else '❌ Not available'}
+            - **GPU**: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None'}
+            """)
             
-            🔑 **Options:**
-            1. **Purchase a license** for unlimited HD downloads
-            2. **Enter your existing license key** if you already purchased
+            # Upload
+            image_input = gr.Image(
+                label="📸 Upload Photo",
+                type="numpy",
+                height=250
+            )
             
-            💡 Visit: https://canadianheadshotpro.carrd.co to purchase
-            📧 Support: bee.tools@zohomailcloud.ca
-            """)
-
-    # ===== CHECK IF MODELS ARE LOADED =====
-    if pipe is None or app is None:
-        raise gr.Error("""
-        ⚠️ AI models are still loading...
-        
-        Please wait a moment and try again.
-        The models take a few minutes to load on first use.
-        
-        If this persists, refresh the page.
-        """)
-
-    # ===== AI GENERATION =====
-    if enable_LCM:
-        pipe.scheduler = diffusers.LCMScheduler.from_config(pipe.scheduler.config)
-        pipe.enable_lora()
-    else:
-        pipe.disable_lora()
-        scheduler_class_name = scheduler.split("-")[0]
-        add_kwargs = {}
-        if len(scheduler.split("-")) > 1:
-            add_kwargs["use_karras_sigmas"] = True
-        if len(scheduler.split("-")) > 2:
-            add_kwargs["algorithm_type"] = "sde-dpmsolver++"
-        scheduler = getattr(diffusers, scheduler_class_name)
-        pipe.scheduler = scheduler.from_config(pipe.scheduler.config, **add_kwargs)
-
-    if face_image_path is None:
-        raise gr.Error("Please upload a face image")
-
-    if prompt is None:
-        prompt = "a person"
-
-    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
-
-    try:
-        face_image = load_image(face_image_path)
-        face_image = resize_img(face_image, max_side=1024)
-        face_image_cv2 = convert_from_image_to_cv2(face_image)
-        height, width, _ = face_image_cv2.shape
-
-        face_info = app.get(face_image_cv2)
-
-        if len(face_info) == 0:
-            raise gr.Error("Unable to detect a face in the image. Please upload a different photo with a clear face.")
-
-        face_info = sorted(
-            face_info,
-            key=lambda x: (x["bbox"][2] - x["bbox"][0]) * x["bbox"][3] - x["bbox"][1],
-        )[-1]
-        face_emb = face_info["embedding"]
-        face_kps = draw_kps(convert_from_cv2_to_image(face_image_cv2), face_info["kps"])
-        img_controlnet = face_image
-
-        if enhance_face_region:
-            control_mask = np.zeros([height, width, 3])
-            x1, y1, x2, y2 = face_info["bbox"]
-            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-            control_mask[y1:y2, x1:x2] = 255
-            control_mask = Image.fromarray(control_mask.astype(np.uint8))
-        else:
-            control_mask = None
-
-        if len(controlnet_selection) > 0 and all(s in controlnet_map and controlnet_map[s] is not None for s in controlnet_selection):
-            controlnet_scales = {
-                "canny": canny_strength,
-                "depth": depth_strength,
-            }
-            pipe.controlnet = MultiControlNetModel(
-                [controlnet_identitynet]
-                + [controlnet_map[s] for s in controlnet_selection]
+            # License
+            license_input = gr.Textbox(
+                label="🔑 License Key",
+                value="HEADSHOT-TEST123456",
+                placeholder="Enter HEADSHOT-TEST123456"
             )
-            control_scales = [float(identitynet_strength_ratio)] + [
-                controlnet_scales[s] for s in controlnet_selection
-            ]
-            control_images = [face_kps] + [
-                controlnet_map_fn[s](img_controlnet).resize((width, height))
-                for s in controlnet_selection
-            ]
-        else:
-            pipe.controlnet = controlnet_identitynet
-            control_scales = float(identitynet_strength_ratio)
-            control_images = face_kps
-
-        generator = torch.Generator(device=device).manual_seed(seed)
-
-        pipe.set_ip_adapter_scale(adapter_strength_ratio)
-        images = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image_embeds=face_emb,
-            image=control_images,
-            control_mask=control_mask,
-            controlnet_conditioning_scale=control_scales,
-            num_inference_steps=num_steps,
-            guidance_scale=guidance_scale,
-            height=height,
-            width=width,
-            generator=generator,
-        ).images
-
-        # ===== 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)
+            # Generate Button
+            generate_btn = gr.Button(
+                "✨ Test GPU & Generate",
+                variant="primary",
+                size="lg"
+            )
             
-            # 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)
-        
-    except Exception as e:
-        print(f"Error during generation: {e}")
-        raise gr.Error(f"Error generating image: {str(e)}")
-
-# ===== GRADIO UI =====
-print("🎨 Building Gradio interface...")
-
-css = """
-:root {
-    --primary: #2563eb;
-    --primary-dark: #1d4ed8;
-    --secondary: #7c3aed;
-    --accent: #06b6d4;
-    --success: #10b981;
-    --warning: #f59e0b;
-    --error: #ef4444;
-    --surface: #ffffff;
-    --surface-alt: #f8fafc;
-    --text-primary: #1e293b;
-    --text-secondary: #64748b;
-    --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;
-    box-shadow: var(--shadow);
-    margin: 20px auto;
-    overflow: hidden;
-    border: 1px solid var(--border);
-}
-.hero-section {
-    background: linear-gradient(135deg, var(--primary) 0%, var(--secondary) 100%);
-    color: white;
-    padding: 40px 0;
-    text-align: center;
-    position: relative;
-    overflow: hidden;
-}
-.hero-section::before {
-    content: '';
-    position: absolute;
-    top: 0;
-    left: 0;
-    right: 0;
-    bottom: 0;
-    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;
-    margin-bottom: 16px;
-    background: linear-gradient(135deg, #ffffff 0%, #e2e8f0 100%);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-    background-clip: text;
-}
-.hero-subtitle {
-    font-size: 1.3em;
-    font-weight: 400;
-    opacity: 0.9;
-    max-width: 600px;
-    margin: 0 auto;
-}
-.upload-area {
-    border: 3px dashed var(--border);
-    border-radius: 20px;
-    padding: 40px 20px;
-    text-align: center;
-    background: var(--surface-alt);
-    transition: all 0.3s ease;
-    cursor: pointer;
-}
-.upload-area:hover {
-    border-color: var(--primary);
-    background: #f0f9ff;
-    transform: translateY(-2px);
-}
-.control-card {
-    background: var(--surface);
-    border-radius: 16px;
-    padding: 24px;
-    margin-bottom: 20px;
-    border: 1px solid var(--border);
-    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;
-    border-radius: 12px;
-    background: linear-gradient(135deg, var(--primary), var(--secondary));
-    display: flex;
-    align-items: center;
-    justify-content: center;
-    margin-right: 12px;
-    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;
-    padding: 30px;
-    border: 1px solid var(--border);
-    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;
-    background: var(--surface-alt);
-    border: 1px solid var(--border);
-    margin-bottom: 20px;
-}
-.success-banner {
-    background: linear-gradient(135deg, var(--success), #059669);
-    color: white;
-    padding: 20px;
-    border-radius: 16px;
-    margin-top: 20px;
-    text-align: center;
-}
-.btn-primary {
-    background: linear-gradient(135deg, var(--primary), var(--primary-dark)) !important;
-    color: white !important;
-    border: none !important;
-    border-radius: 12px !important;
-    padding: 16px 32px !important;
-    font-weight: 600 !important;
-    font-size: 1.1em !important;
-    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;
-    padding: 4px 12px;
-    border-radius: 20px;
-    font-size: 0.85em;
-    font-weight: 600;
-}
-.tips-card {
-    background: linear-gradient(135deg, #fef3c7, #f59e0b);
-    border: none;
-    border-radius: 16px;
-    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 {
-        font-size: 2em;
-    }
-    .hero-subtitle {
-        font-size: 1.1em;
-    }
-    .control-card {
-        padding: 16px;
-    }
-}
-.loading-overlay {
-    position: fixed;
-    top: 0;
-    left: 0;
-    right: 0;
-    bottom: 0;
-    background: rgba(255, 255, 255, 0.9);
-    display: flex;
-    align-items: center;
-    justify-content: center;
-    z-index: 1000;
-}
-"""
-
-with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Pro AI Headshot Generator") as demo:
-    
-    # Loading indicator
-    loading_html = gr.HTML("""
-    <div class="loading-overlay">
-        <div style="text-align: center;">
-            <h2>🚀 Loading AI Headshot Generator...</h2>
-            <p>This may take 2-3 minutes on first load. Please be patient.</p>
-            <p>Models are downloading and initializing...</p>
-        </div>
-    </div>
-    """, visible=True)
-    
-    # Main Container
-    with gr.Column(elem_classes="main-container") as main_container:
-        
-        # Hero Section
-        with gr.Column(elem_classes="hero-section"):
-            gr.HTML("""
-            <div style="position: relative; z-index: 2;">
-                <h1 class="hero-title">🎯 Pro AI Headshot Generator</h1>
-                <p class="hero-subtitle">Transform any selfie into professional headshots in seconds. Perfect for LinkedIn, corporate profiles, and professional portfolios.</p>
-            </div>
-            """)
+            # Status
+            status_output = gr.Textbox(
+                label="Status",
+                interactive=False
+            )
         
-        # Main Content
-        with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            # Result
+            result_image = gr.Image(
+                label="🖼️ Result",
+                height=400
+            )
             
-            # Left Column - Controls
-            with gr.Column(scale=1, min_width=400):
-                
-                # Upload Section
-                with gr.Column(elem_classes="control-card"):
-                    gr.HTML("""
-                    <div class="control-header">
-                        <div class="control-icon">📸</div>
-                        <h3 class="control-title">Upload Your Photo</h3>
-                    </div>
-                    """)
-                    gr.HTML("""
-                    <p style="color: var(--text-secondary); margin-bottom: 20px; font-size: 0.95em;">
-                        For best results, use a clear, well-lit photo where your face is clearly visible.
-                    </p>
-                    """)
-                    face_file = gr.Image(
-                        label="",
-                        type="filepath",
-                        height=200,
-                        show_label=False,
-                        elem_classes="upload-area"
-                    )
-                
-                # Access Options Section
-                with gr.Column(elem_classes="control-card"):
-                    gr.HTML("""
-                    <div class="control-header">
-                        <div class="control-icon">🔑</div>
-                        <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 license key (or leave blank for free trial)",
-                        show_label=False,
-                        value="HEADSHOT-TEST123456",
-                        info=f"💡 You get {MAX_FREE_TRIALS} free generations. Purchase license for HD downloads without watermark."
-                    )
-                    
-                    gr.HTML(f"""
-                    <div style="font-size: 0.85em; color: var(--text-secondary); margin-top: 8px;">
-                        <strong>Free Trial:</strong> {MAX_FREE_TRIALS} watermarked previews<br>
-                        <strong>Premium License:</strong> Unlimited HD downloads, no watermark<br>
-                        <strong>Test Key:</strong> HEADSHOT-TEST123456<br>
-                        <a href="https://canadianheadshotpro.carrd.co" target="_blank" style="color: var(--primary); font-weight: 600;">👉 Click here to purchase license</a>
-                    </div>
-                    """)
-                
-                # Description Section
-                with gr.Column(elem_classes="control-card"):
-                    gr.HTML("""
-                    <div class="control-header">
-                        <div class="control-icon">✍️</div>
-                        <h3 class="control-title">Describe Your Look</h3>
-                    </div>
-                    """)
-                    prompt = gr.Textbox(
-                        label="",
-                        placeholder="Describe how you want to appear...",
-                        value="modern professional headshot, creative director style, soft natural lighting, authentic expression, contemporary business portrait, premium quality photo",
-                        show_label=False,
-                        lines=3
-                    )
-                    gr.HTML("""
-                    <div style="font-size: 0.85em; color: var(--text-secondary); margin-top: 8px;">
-                        💡 Examples: "professional business headshot", "friendly corporate portrait", "creative director style"
-                    </div>
-                    """)
-                
-                # Style Selection
-                with gr.Column(elem_classes="control-card"):
-                    gr.HTML("""
-                    <div class="control-header">
-                        <div class="control-icon">🎨</div>
-                        <h3 class="control-title">Style Options</h3>
-                    </div>
-                    """)
-                    style = gr.Dropdown(
-                        label="Style Theme",
-                        choices=["No Style"] + STYLE_NAMES,
-                        value="No Style",
-                        info="'No Style' recommended for natural professional results"
-                    )
-                
-                # Quality Settings
-                with gr.Column(elem_classes="control-card"):
-                    gr.HTML("""
-                    <div class="control-header">
-                        <div class="control-icon">⚙️</div>
-                        <h3 class="control-title">Quality Settings</h3>
-                    </div>
-                    """)
-                    
-                    with gr.Row():
-                        identitynet_strength_ratio = gr.Slider(
-                            label="Face Similarity",
-                            minimum=0.5,
-                            maximum=1.2,
-                            step=0.05,
-                            value=0.80,
-                            info="How closely the headshot resembles your photo"
-                        )
-                    
-                    with gr.Row():
-                        adapter_strength_ratio = gr.Slider(
-                            label="Detail Quality",
-                            minimum=0.3,
-                            maximum=1.2,
-                            step=0.05,
-                            value=0.55,
-                            info="Level of detail in the final image"
-                        )
-                    
-                    enable_LCM = gr.Checkbox(
-                        label="Enable Fast Generation Mode",
-                        value=False,
-                        info="Faster results with slightly lower quality"
-                    )
-                
-                # Tips Card
-                with gr.Column(elem_classes="tips-card"):
-                    gr.HTML("""
-                    <div class="tips-header">
-                        <div class="tips-icon">💡</div>
-                        <h4 style="margin: 0; color: #92400e;">Pro Tips for Best Results</h4>
-                    </div>
-                    <ul style="margin: 0; color: #92400e; font-size: 0.9em;">
-                        <li>Use clear, well-lit face photos</li>
-                        <li>Face should be visible and not too small</li>
-                        <li>Avoid blurry or dark images</li>
-                        <li>Single person in photo works best</li>
-                    </ul>
-                    """)
-                
-                # Generate Button
-                submit = gr.Button(
-                    "✨ Generate Professional Headshot", 
-                    variant="primary", 
-                    size="lg",
-                    elem_classes="btn-primary",
-                    scale=1
-                )
+            gr.Markdown("""
+            ### 🎯 Next Steps
+            1. ✅ This confirms ZeroGPU works
+            2. 🔧 Add AI models gradually
+            3. 🚀 Deploy full headshot generator
             
-            # Right Column - Results
-            with gr.Column(scale=1, min_width=500):
-                with gr.Column(elem_classes="result-card"):
-                    gr.HTML("""
-                    <div class="result-header">
-                        <h2 class="result-title">Your Professional Headshot</h2>
-                        <p class="result-subtitle">Your AI-generated headshot will appear here. Download as high-quality PNG for professional use.</p>
-                    </div>
-                    """)
-                    
-                    # Image Display
-                    gallery = gr.Image(
-                        label="",
-                        height=400,
-                        show_download_button=True,
-                        show_label=False,
-                        type="filepath",
-                        elem_classes="image-container"
-                    )
-                    
-                    # Success Message
-                    success_msg = gr.HTML("""
-                    <div class="success-banner" style="display: none;">
-                        <h4 style="margin: 0 0 8px 0;">✅ Success! Your Professional Headshot is Ready</h4>
-                        <p style="margin: 0; opacity: 0.9;">Download your high-quality PNG file for LinkedIn, professional profiles, or portfolios.</p>
-                    </div>
-                    """)
-                    
-                    # Progress Information
-                    progress_info = gr.HTML("""
-                    <div class="progress-container">
-                        <div style="font-size: 0.9em; color: var(--text-secondary);">
-                            ⏱️ Generation takes 20-30 seconds
-                        </div>
-                    </div>
-                    """)
-
-    # Hidden technical parameters
-    negative_prompt = gr.Textbox(
-        value="(lowres, low quality, worst quality:1.2), (text:1.2), watermark, (frame:1.2), deformed, ugly, deformed eyes, blur, out of focus, blurry, deformed cat, deformed, photo, anthropomorphic cat, monochrome, pet collar, gun, weapon, blue, 3d, drones, drone, buildings in background, green",
-        visible=False
-    )
-    num_steps = gr.Slider(
-        minimum=5,
-        maximum=100,
-        step=1,
-        value=30,
-        label="Number of steps",
-        visible=False
-    )
-    guidance_scale = gr.Slider(
-        minimum=0.1,
-        maximum=20.0,
-        step=0.1,
-        value=5.0,
-        label="Guidance scale",
-        visible=False
-    )
-    seed = gr.Slider(
-        minimum=0,
-        maximum=MAX_SEED,
-        step=1,
-        value=42,
-        label="Seed",
-        visible=False
-    )
-    scheduler = gr.Dropdown(
-        value="EulerDiscreteScheduler",
-        choices=["EulerDiscreteScheduler", "EulerAncestralDiscreteScheduler", "DPMSolverMultistepScheduler"],
-        visible=False
-    )
-    randomize_seed = gr.Checkbox(value=True, visible=False)
-    enhance_face_region = gr.Checkbox(value=True, visible=False)
-    controlnet_selection = gr.CheckboxGroup(
-        choices=["canny", "depth"],
-        value=["depth"],
-        label="Controlnet",
-        visible=False
-    )
-    canny_strength = gr.Slider(
-        minimum=0,
-        maximum=1.5,
-        step=0.01,
-        value=0.4,
-        label="Canny strength",
-        visible=False
-    )
-    depth_strength = gr.Slider(
-        minimum=0,
-        maximum=1.5,
-        step=0.01,
-        value=0.4,
-        label="Depth strength",
-        visible=False
-    )
-
-    def show_success():
-        return """
-        <div class="success-banner">
-            <h4 style="margin: 0 0 8px 0;">✅ Success! Your Professional Headshot is Ready</h4>
-            <p style="margin: 0; opacity: 0.9;">Download your high-quality PNG file for LinkedIn, professional profiles, or portfolios.</p>
-        </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>
-        """
-    
-    def hide_loading():
-        """Hide loading overlay once app is ready"""
-        return gr.update(visible=False)
-
-    # Hide loading overlay after 5 seconds
-    import threading
-    def delayed_hide():
-        time.sleep(5)
-        loading_html.value = ""
+            [🔗 Purchase License](https://canadianheadshotpro.carrd.co)
+            """)
     
-    hide_thread = threading.Thread(target=delayed_hide)
-    hide_thread.daemon = True
-    hide_thread.start()
-
-    submit.click(
-        fn=generate_image,
-        inputs=[
-            face_file,
-            license_input,
-            prompt,
-            negative_prompt,
-            style,
-            num_steps,
-            identitynet_strength_ratio,
-            adapter_strength_ratio,
-            canny_strength,
-            depth_strength,
-            controlnet_selection,
-            guidance_scale,
-            seed,
-            scheduler,
-            enable_LCM,
-            enhance_face_region,
-        ],
-        outputs=[gallery, success_msg]
-    ).then(
-        fn=show_success,
-        outputs=success_msg
-    )
-
-    # Update trial display when license input changes
-    license_input.change(
-        fn=update_trial_display,
-        inputs=[license_input],
-        outputs=[trial_status]
+    # Connect
+    generate_btn.click(
+        fn=generate_headshot,
+        inputs=[image_input, license_input],
+        outputs=[result_image, status_output]
     )
-
-    enable_LCM.input(
-        fn=toggle_lcm_ui,
-        inputs=[enable_LCM],
-        outputs=[num_steps, guidance_scale],
-        queue=False,
+    
+    # Auto-test on load
+    demo.load(
+        fn=lambda: f"��� ZeroGPU ready with {torch.cuda.get_device_name(0)}",
+        outputs=[status_output]
     )
 
-print("✅ App initialization complete!")
-print(f"🔑 Test license key: HEADSHOT-TEST123456")
-print(f"🎮 Free trials: {MAX_FREE_TRIALS} per user")
-print("🚀 Ready to launch!")
-
 if __name__ == "__main__":
-    demo.queue(max_size=3)
-    demo.launch(
-        share=True,
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        debug=True
-    )
+    demo.launch(share=True)