Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |
| ) | |