Update app.py

app.py

@@ -4,11 +4,26 @@ os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
-import torch
 
-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)}")
+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"}
@@ -19,11 +34,11 @@ def verify_license(key):
 
 # ===== SIMPLE GENERATOR =====
 def generate_headshot(image, license_key):
-    """Placeholder that shows GPU is working"""
+    """Generate placeholder headshot"""
     
     # Verify license
     if not verify_license(license_key):
-        return np.zeros((512, 512, 3), dtype=np.uint8), "❌ Invalid license"
+        return np.zeros((512, 512, 3), dtype=np.uint8), "❌ Invalid license. Use: HEADSHOT-TEST123456"
     
     # Create result image
     if image is not None:
@@ -32,35 +47,33 @@ def generate_headshot(image, license_key):
     else:
         result = np.ones((512, 512, 3), dtype=np.uint8) * 240
     
-    # Add GPU status
+    # Add info
     img = Image.fromarray(result)
     draw = ImageDraw.Draw(img)
     
-    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))
+    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), "✅ Success! GPU is working."
+    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("""
+    gr.Markdown(f"""
     # 🎯 AI Headshot Generator
-    *Powered by ZeroGPU - PyTorch ready*
+    *Running on: {'🚀 GPU' if HAS_CUDA else '⚡ CPU'}*
     """)
     
     with gr.Row():
         with gr.Column(scale=1):
-            # GPU Status
+            # Hardware 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'}
+            ### ⚙️ Hardware Status
+            - **PyTorch**: {'✅ Available' if HAS_TORCH else '❌ Not available'}
+            - **CUDA**: {'✅ Available' if HAS_CUDA else '⚠️ CPU mode'}
             """)
             
             # Upload
@@ -79,7 +92,7 @@ with gr.Blocks(title="AI Headshot Generator", theme=gr.themes.Soft()) as demo:
             
             # Generate Button
             generate_btn = gr.Button(
-                "✨ Test GPU & Generate",
+                "✨ Generate Headshot",
                 variant="primary",
                 size="lg"
             )
@@ -89,35 +102,36 @@ with gr.Blocks(title="AI Headshot Generator", theme=gr.themes.Soft()) as demo:
                 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="🖼️ Result",
+                label="🖼️ Your Headshot",
                 height=400
             )
             
             gr.Markdown("""
-            ### 🎯 Next Steps
-            1. ✅ This confirms ZeroGPU works
-            2. 🔧 Add AI models gradually
-            3. 🚀 Deploy full headshot generator
+            ### 📥 Download Options
+            - **Free Trial**: Watermarked preview
+            - **Premium**: HD quality, no watermark
+            - **Professional**: Multiple styles
             
-            [🔗 Purchase License](https://canadianheadshotpro.carrd.co)
+            ⏱️ *Generation time: 10-20 seconds*
             """)
-    
+
     # Connect
     generate_btn.click(
         fn=generate_headshot,
         inputs=[image_input, license_input],
         outputs=[result_image, status_output]
     )
-    
-    # Auto-test on load
-    demo.load(
-        fn=lambda: f"✅ ZeroGPU ready with {torch.cuda.get_device_name(0)}",
-        outputs=[status_output]
-    )
 
 if __name__ == "__main__":
     demo.launch(share=True)