app.py

"""

Shelter Flyer Generator - Main Application

A Gradio-based web app for animal shelters to generate adoption flyers.



Compatible with both local deployment and Hugging Face Spaces.

"""

import os
import sys
import gradio as gr
from PIL import Image
from pathlib import Path
from datetime import datetime

# Load environment variables (for local .env support)
# Priority: Environment variables > .env file
try:
    from dotenv import load_dotenv
    load_dotenv()  # This loads .env if present, but doesn't override existing env vars
except ImportError:
    print("[INFO] python-dotenv not installed. Using environment variables only.")

# Import core modules
try:
    from core.background_removal import remove_background
    from core.compositing import composite_animal_on_scene
    from core.bio_generator import generate_bio
    from core.flyer_builder import build_flyer
    print("[OK] All core modules loaded successfully.")
except ImportError as e:
    print(f"[ERROR] Failed to import core modules: {e}")
    print("Please ensure all core modules are present in the 'core/' directory.")
    sys.exit(1)

# ============================================================================
# ROBUST TOKEN MANAGEMENT
# ============================================================================
def check_hf_token():
    """

    Check for HF_TOKEN with proper priority:

    1. Environment variable (for HF Spaces)

    2. .env file (for local development)

    

    Returns:

        str or None: The token if found, None otherwise

    """
    token = os.getenv("HF_TOKEN")
    
    if token and token != "hf_xxxxxxxxxxxxxxxxxxxxx":
        print("✅ Token detected")
        return token
    else:
        print("❌ Token not found")
        print("[WARN] Bio generation will use fallback template mode.")
        print("[INFO] To use AI bio generation, set HF_TOKEN environment variable.")
        print("[INFO] Get your token at: https://huggingface.co/settings/tokens")
        return None

# Check token at startup
HF_TOKEN = check_hf_token()

# ============================================================================
# PATH CONSISTENCY (pathlib for cross-platform compatibility)
# ============================================================================
BASE_DIR = Path(__file__).parent.resolve()
BACKGROUNDS_DIR = BASE_DIR / "backgrounds"
OUTPUTS_DIR = BASE_DIR / "outputs"

# Ensure outputs directory exists
OUTPUTS_DIR.mkdir(parents=True, exist_ok=True)

# Pre-defined personality traits
PERSONALITY_TRAITS = [
    "Good with kids",
    "Good with other dogs",
    "Good with cats",
    "High energy",
    "Calm / Relaxed",
    "Playful",
    "Loyal",
    "Trained / Knows commands",
    "Loves cuddles",
    "Independent",
    "Gentle",
    "Protective",
    "Loves walks",
    "House-trained",
    "Special needs"
]

# Available background scenes
BACKGROUND_SCENES = {
    "Living Room": "living_room.jpg",
    "Cozy Home": "cozy_home.jpg",
    "Garden/Yard": "garden.jpg",
    "Forest Park": "park.jpg",
    "Beach Sunset": "beach.jpg"
}


# ============================================================================
# MAIN PIPELINE WITH FILE DOWNLOAD SUPPORT
# ============================================================================
def generate_flyer_pipeline(

    animal_photo,

    scene_choice,

    custom_scene,

    animal_name,

    animal_type,

    breed,

    age,

    traits,

    additional_notes,

    shelter_name,

    shelter_contact

):
    """

    Main pipeline that orchestrates the flyer generation process.

    

    Returns:

        enhanced_photo: PIL.Image - The composited animal on scene

        generated_bio: str - The AI-generated bio

        flyer_file_path: str - Path to downloadable flyer file

        status: str - Status message

    """
    
    try:
        # Validate inputs
        if animal_photo is None:
            return None, None, None, "❌ Please upload an animal photo."
        
        if not animal_name or not animal_type or not breed or not age:
            return None, None, None, "❌ Please fill in all animal information fields."
        
        status_msg = "🚀 Starting flyer generation...\n\n"
        
        # Step 1: Select scene
        if custom_scene is not None:
            scene_image = custom_scene
            status_msg += "✅ Using custom background scene\n"
        elif scene_choice and scene_choice in BACKGROUND_SCENES:
            scene_path = BACKGROUNDS_DIR / BACKGROUND_SCENES[scene_choice]
            if not scene_path.exists():
                return None, None, None, f"❌ Background file not found: {scene_path}"
            scene_image = Image.open(scene_path)
            status_msg += f"✅ Using {scene_choice} background\n"
        else:
            return None, None, None, "❌ Please select a background scene or upload a custom one."
        
        # Step 2: Remove background from animal photo
        status_msg += "🔄 Removing background from animal photo...\n"
        try:
            animal_cutout = remove_background(animal_photo)
            status_msg += "✅ Background removed successfully\n"
        except Exception as e:
            status_msg += f"❌ Background removal failed: {str(e)}\n"
            return None, None, None, status_msg
        
        # Step 3: Composite animal onto scene
        status_msg += "🔄 Compositing animal onto scene...\n"
        try:
            enhanced_photo = composite_animal_on_scene(
                animal_rgba=animal_cutout,
                scene=scene_image,
                position="center",
                scale_factor=0.6
            )
            status_msg += "✅ Enhanced photo created\n"
        except Exception as e:
            status_msg += f"❌ Compositing failed: {str(e)}\n"
            return None, None, None, status_msg
        
        # Step 4: Generate bio with AI error handling & fallback
        status_msg += "🔄 Generating bio...\n"
        try:
            bio_text = generate_bio(
                name=animal_name,
                animal_type=animal_type,
                breed=breed,
                age=age,
                traits=traits if traits else [],
                additional_notes=additional_notes
            )
            
            # Check if fallback was used
            if "Meet" in bio_text and "looking for a forever home" in bio_text:
                status_msg += "✅ Bio generated (using template fallback)\n"
            else:
                status_msg += "✅ Bio generated with AI\n"
                
        except Exception as e:
            status_msg += f"⚠️ Bio generation encountered an issue: {str(e)}\n"
            status_msg += "📝 Using template-based bio as fallback...\n"
            
            # Manual fallback if bio_generator fails completely
            bio_text = f"Meet {animal_name}! This {age} {breed} {animal_type.lower()} is looking for a forever home. "
            if traits:
                bio_text += f"{animal_name} is {', '.join(traits[:3]).lower()}. "
            bio_text += f"Come meet {animal_name} at our shelter today!"
            status_msg += "✅ Template bio created\n"
        
        # Step 5: Build final flyer
        status_msg += "🔄 Building final flyer...\n"
        try:
            flyer_image = build_flyer(
                photo=enhanced_photo,
                name=animal_name,
                animal_type=animal_type,
                breed=breed,
                age=age,
                bio=bio_text,
                shelter_name=shelter_name if shelter_name else "Your Local Animal Shelter",
                shelter_contact=shelter_contact if shelter_contact else "",
                template="template_1"
            )
            status_msg += "✅ Flyer created successfully\n"
        except Exception as e:
            status_msg += f"❌ Flyer building failed: {str(e)}\n"
            return None, None, None, status_msg
        
        # Step 6: Save flyer to file for download
        status_msg += "💾 Saving flyer...\n"
        try:
            # Create unique filename with timestamp
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            safe_name = "".join(c for c in animal_name if c.isalnum() or c in (' ', '-', '_'))
            safe_name = safe_name.replace(' ', '_')
            flyer_filename = f"flyer_{safe_name}_{timestamp}.jpg"
            flyer_path = OUTPUTS_DIR / flyer_filename
            
            # Save with high quality
            flyer_image.save(flyer_path, "JPEG", quality=95, optimize=True)
            
            status_msg += f"✅ Flyer saved: {flyer_filename}\n"
            status_msg += "\n🎉 Generation complete! You can now download your flyer.\n"
            
            # Return the path as a string (Gradio will handle the file download)
            return enhanced_photo, bio_text, str(flyer_path), status_msg
            
        except Exception as e:
            status_msg += f"❌ Failed to save flyer: {str(e)}\n"
            return enhanced_photo, bio_text, None, status_msg
        
    except Exception as e:
        import traceback
        error_msg = "❌ Unexpected error occurred:\n\n"
        error_msg += f"Error: {str(e)}\n\n"
        error_msg += "Stack trace:\n"
        error_msg += traceback.format_exc()
        return None, None, None, error_msg


def regenerate_bio_only(

    animal_name,

    animal_type,

    breed,

    age,

    traits,

    additional_notes

):
    """

    Regenerates just the bio without reprocessing images.

    Includes proper error handling and fallback.

    """
    
    try:
        if not animal_name or not animal_type:
            return "❌ Please provide at least name and animal type.", "Error"
        
        bio_text = generate_bio(
            name=animal_name,
            animal_type=animal_type,
            breed=breed,
            age=age,
            traits=traits if traits else [],
            additional_notes=additional_notes
        )
        
        # Check if AI or fallback was used
        if "Meet" in bio_text and "looking for a forever home" in bio_text:
            status = "✅ Bio regenerated (template mode)"
        else:
            status = "✅ Bio regenerated with AI"
        
        return bio_text, status
        
    except Exception as e:
        error_msg = f"❌ Error generating bio: {str(e)}\n"
        error_msg += "Using fallback template..."
        
        # Fallback bio
        bio_text = f"Meet {animal_name}! This {breed} {animal_type.lower()} is looking for a loving home. "
        if traits:
            bio_text += f"{animal_name} is {', '.join(traits[:2]).lower()}. "
        bio_text += f"Visit our shelter to meet {animal_name}!"
        
        return bio_text, "⚠️ Used fallback template (API unavailable)"


# ============================================================================
# GRADIO 5.0+ COMPATIBLE UI
# ============================================================================
with gr.Blocks(
    title="🐾 Shelter Flyer Generator",
    theme=gr.themes.Soft()
) as app:
    
    gr.Markdown("""

    # 🐾 Shelter Flyer Generator

    

    Create beautiful adoption flyers for shelter animals with AI-enhanced photos and compelling bios.

    

    **How it works:**

    1. Upload a photo of the animal

    2. Choose a background scene

    3. Fill in the animal's information

    4. Generate a professional flyer!

    """)
    
    with gr.Tabs():
        # ====================================================================
        # TAB 1: Upload & Describe
        # ====================================================================
        with gr.Tab("📸 Upload & Describe"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Animal Photo")
                    animal_photo = gr.Image(
                        type="pil",
                        label="Upload Animal Photo",
                        height=400
                    )
                    
                    gr.Markdown("### Background Scene")
                    scene_choice = gr.Radio(
                        choices=list(BACKGROUND_SCENES.keys()),
                        label="Choose a Background Scene",
                        value="Living Room"
                    )
                    
                    custom_scene = gr.Image(
                        type="pil",
                        label="Or Upload Your Own Scene (optional)",
                        height=200
                    )
                
                with gr.Column():
                    gr.Markdown("### Animal Information")
                    animal_name = gr.Textbox(
                        label="Animal Name",
                        placeholder="e.g., Bella"
                    )
                    
                    animal_type = gr.Radio(
                        choices=["Dog", "Cat", "Other"],
                        label="Animal Type",
                        value="Dog"
                    )
                    
                    with gr.Row():
                        breed = gr.Textbox(
                            label="Breed",
                            placeholder="e.g., Golden Retriever"
                        )
                        age = gr.Textbox(
                            label="Age",
                            placeholder="e.g., 2 years"
                        )
                    
                    gr.Markdown("### Personality Traits")
                    traits = gr.CheckboxGroup(
                        choices=PERSONALITY_TRAITS,
                        label="Select all that apply"
                    )
                    
                    additional_notes = gr.Textbox(
                        lines=3,
                        label="Additional Notes",
                        placeholder="Any special information about this animal..."
                    )
                    
                    gr.Markdown("### Shelter Information")
                    shelter_name = gr.Textbox(
                        label="Shelter Name",
                        placeholder="e.g., Happy Paws Animal Shelter",
                        value="Your Local Animal Shelter"
                    )
                    
                    shelter_contact = gr.Textbox(
                        label="Contact Info",
                        placeholder="e.g., (555) 123-4567 | www.shelter.org"
                    )
            
            with gr.Row():
                generate_btn = gr.Button(
                    "🐾 Generate Flyer",
                    variant="primary",
                    size="lg"
                )
            
            status_output = gr.Textbox(
                label="Status",
                lines=10,
                interactive=False
            )
        
        # ====================================================================
        # TAB 2: Preview & Edit
        # ====================================================================
        with gr.Tab("👁️ Preview & Edit"):
            gr.Markdown("### Enhanced Photo")
            enhanced_photo_output = gr.Image(
                label="Enhanced Photo Preview",
                type="pil"
            )
            
            gr.Markdown("### Generated Bio")
            with gr.Row():
                bio_output = gr.Textbox(
                    lines=8,
                    label="Generated Bio (Editable)",
                    interactive=True
                )
            
            with gr.Row():
                regenerate_bio_btn = gr.Button("🔄 Regenerate Bio")
                bio_status = gr.Textbox(
                    label="Bio Status",
                    lines=1,
                    interactive=False
                )
        
        # ====================================================================
        # TAB 3: Download (UPDATED FOR FILE DOWNLOADS)
        # ====================================================================
        with gr.Tab("⬇️ Download"):
            gr.Markdown("""

            ### Download Your Flyer

            

            Once your flyer is generated, you can download it here.

            """)
            
            # Main flyer download with file output
            download_flyer = gr.File(
                label="📥 Download Final Flyer (JPG)",
                interactive=False
            )
            
            gr.Markdown("""

            ---

            **Tip:** The downloaded JPG is high-quality and ready for:

            - Social media posts

            - Email campaigns  

            - Printing on flyers

            - Posting on your shelter website

            """)
    
    # ========================================================================
    # EVENT HANDLERS
    # ========================================================================
    generate_btn.click(
        fn=generate_flyer_pipeline,
        inputs=[
            animal_photo,
            scene_choice,
            custom_scene,
            animal_name,
            animal_type,
            breed,
            age,
            traits,
            additional_notes,
            shelter_name,
            shelter_contact
        ],
        outputs=[
            enhanced_photo_output,
            bio_output,
            download_flyer,  # Now outputs file path for download
            status_output
        ]
    )
    
    regenerate_bio_btn.click(
        fn=regenerate_bio_only,
        inputs=[
            animal_name,
            animal_type,
            breed,
            age,
            traits,
            additional_notes
        ],
        outputs=[
            bio_output,
            bio_status
        ]
    )
    
    gr.Markdown("""

    ---

    ### About

    

    This app uses:

    - **Background Removal**: RMBG-2.0 (Hugging Face)

    - **Bio Generation**: Zephyr-7B (Hugging Face Inference API)

    - **Image Processing**: Pillow

    

    Built for animal shelters to create compelling adoption flyers. 🐾

    

    **Note:** If HF_TOKEN is not configured, bio generation will use template mode.

    """)


# ============================================================================
# LAUNCH CONFIGURATION
# ============================================================================
if __name__ == "__main__":
    print("\n" + "="*60)
    print("🐾 Shelter Flyer Generator")
    print("="*60)
    print(f"📁 Base directory: {BASE_DIR}")
    print(f"📁 Backgrounds: {BACKGROUNDS_DIR}")
    print(f"📁 Outputs: {OUTPUTS_DIR}")
    print(f"🔑 HF Token: {'✅ Configured' if HF_TOKEN else '❌ Not found (using fallback mode)'}")
    print("="*60 + "\n")
    
    # Launch with proper settings for both local and HF Spaces
    app.launch(
        server_name="0.0.0.0",  # Allows external access (needed for HF Spaces)
        server_port=7860,        # Standard port
        share=False              # Set to True for temporary public link
    )