app.py

import gradio as gr
import yaml
import torch
from PIL import Image
from rembg import remove
from rate_limiter import RateLimiter
from replicate_handler import ReplicateHandler


# ------------------------------------------
# Load Configuration
# ------------------------------------------

with open("config.yaml", "r") as f:
    config = yaml.safe_load(f)

# Initialize modules
rate_limiter = RateLimiter(
    session_file=config["rate_limit"]["session_file"],
    daily_limit=config["rate_limit"]["default_daily_limit"],
    dev_daily_limit=config["rate_limit"]["dev_daily_limit"]
)

replicate_handler = ReplicateHandler(
    model=config["replicate"]["model"],
    default_settings=config["replicate"]["default_settings"]
)

# ------------------------------------------
# Core Functions
# ------------------------------------------

def remove_background(img: Image.Image) -> Image.Image:
    """Remove background from image"""
    if img is None:
        raise gr.Error("Please upload an image first")
    return remove(img)


def generate_headshot(
    img: Image.Image,
    headshot_style: str,
    use_advanced: bool,
    cfg: float,
    steps: int,
    denoise: float,
    request: gr.Request
) -> Image.Image:
    """Generate professional headshot with selected style"""
    
    # Check rate limit
    allowed, remaining, reset_time = rate_limiter.check_limit(request)
    if not allowed:
        raise gr.Error(rate_limiter.get_limit_message(remaining, reset_time))
    
    if img is None:
        raise gr.Error("Please upload an image first")
    
    # Get prompt from config
    style_config = config["headshots"][headshot_style]
    prompt = style_config["prompt"]
    negative_prompt = style_config["negative"]
    
    # Prepare settings
    custom_settings = None
    if use_advanced:
        custom_settings = {
            "cfg": cfg,
            "steps": steps,
            "denoise": denoise
        }
    
    # Generate
    try:
        result = replicate_handler.generate(
            input_image=img,
            prompt=prompt,
            negative_prompt=negative_prompt,
            custom_settings=custom_settings
        )
        
        # Increment usage
        rate_limiter.increment(request)
        
        return result
        
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")


def generate_scene(
    img: Image.Image,
    scene_style: str,
    use_advanced: bool,
    cfg: float,
    steps: int,
    denoise: float,
    request: gr.Request
) -> Image.Image:
    """Generate scene change with selected style"""
    
    # Check rate limit
    allowed, remaining, reset_time = rate_limiter.check_limit(request)
    if not allowed:
        raise gr.Error(rate_limiter.get_limit_message(remaining, reset_time))
    
    if img is None:
        raise gr.Error("Please upload an image first")
    
    # Get prompt from config
    style_config = config["scenes"][scene_style]
    prompt = style_config["prompt"]
    negative_prompt = style_config["negative"]
    
    # Prepare settings
    custom_settings = None
    if use_advanced:
        custom_settings = {
            "cfg": cfg,
            "steps": steps,
            "denoise": denoise
        }
    
    # Generate
    try:
        result = replicate_handler.generate(
            input_image=img,
            prompt=prompt,
            negative_prompt=negative_prompt,
            custom_settings=custom_settings
        )
        
        # Increment usage
        rate_limiter.increment(request)
        
        return result
        
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")


def get_rate_limit_status(request: gr.Request) -> str:
    """Get current rate limit status"""
    allowed, remaining, reset_time = rate_limiter.check_limit(request)
    return rate_limiter.get_limit_message(remaining, reset_time)


# ------------------------------------------
# Gradio Interface
# ------------------------------------------

with gr.Blocks(theme=gr.themes.Soft(), title="FaceForge AI") as demo:
    gr.Markdown(
        """
        # 🎨 FaceForge AI
        ### AI-Powered Professional Photo Generator
        Upload your photo and transform it into professional headshots or place yourself in different scenes.
        """
    )
    
    # Rate limit status
    rate_status = gr.Textbox(
        label="📊 Usage Status",
        interactive=False,
        value="Loading..."
    )
    
    with gr.Row():
        # Left Column - Input
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="📷 Upload Your Photo")
            
            gr.Markdown("### ⚙️ Advanced Settings (Optional)")
            
            use_advanced = gr.Checkbox(
                label="Enable Advanced Settings",
                value=False
            )
            
            with gr.Group(visible=True) as advanced_group:
                cfg_slider = gr.Slider(
                    label="🎛️ CFG Scale",
                    minimum=config["replicate"]["advanced_ranges"]["cfg"]["min"],
                    maximum=config["replicate"]["advanced_ranges"]["cfg"]["max"],
                    value=config["replicate"]["default_settings"]["cfg"],
                    step=config["replicate"]["advanced_ranges"]["cfg"]["step"]
                )
                
                steps_slider = gr.Slider(
                    label="🔄 Steps",
                    minimum=config["replicate"]["advanced_ranges"]["steps"]["min"],
                    maximum=config["replicate"]["advanced_ranges"]["steps"]["max"],
                    value=config["replicate"]["default_settings"]["steps"],
                    step=config["replicate"]["advanced_ranges"]["steps"]["step"]
                )
                
                denoise_slider = gr.Slider(
                    label="✨ Denoise",
                    minimum=config["replicate"]["advanced_ranges"]["denoise"]["min"],
                    maximum=config["replicate"]["advanced_ranges"]["denoise"]["max"],
                    value=config["replicate"]["default_settings"]["denoise"],
                    step=config["replicate"]["advanced_ranges"]["denoise"]["step"]
                )
            
            # Toggle visibility of advanced settings
            use_advanced.change(
                fn=lambda x: gr.update(visible=x),
                inputs=[use_advanced],
                outputs=[advanced_group]
            )
        
        # Right Column - Outputs
        with gr.Column(scale=1):
            
            # Background Remover Section
            gr.Markdown("### 🎭 Background Remover")
            output_bg_removed = gr.Image(label="Background Removed", type="pil")
            btn_bg_remove = gr.Button("🗑️ Remove Background", variant="secondary")
            
            gr.Markdown("---")
            
            # Professional Headshots Section
            gr.Markdown("### 💼 Professional Headshots")
            headshot_style = gr.Dropdown(
                label="Choose Headshot Style",
                choices=list(config["headshots"].keys()),
                value=list(config["headshots"].keys())[0]
            )
            output_headshot = gr.Image(label="Professional Headshot", type="pil")
            btn_headshot = gr.Button("📸 Generate Headshot", variant="primary")
            
            gr.Markdown("---")
            
            # Scene Changer Section
            gr.Markdown("### 🌍 Scene Changer")
            scene_style = gr.Dropdown(
                label="Choose Scene",
                choices=list(config["scenes"].keys()),
                value=list(config["scenes"].keys())[0]
            )
            output_scene = gr.Image(label="Scene Changed", type="pil")
            btn_scene = gr.Button("🎬 Change Scene", variant="primary")
    
    # Button Actions
    btn_bg_remove.click(
        fn=remove_background,
        inputs=[input_image],
        outputs=[output_bg_removed]
    )
    
    btn_headshot.click(
        fn=generate_headshot,
        inputs=[
            input_image,
            headshot_style,
            use_advanced,
            cfg_slider,
            steps_slider,
            denoise_slider
        ],
        outputs=[output_headshot]
    )
    
    btn_scene.click(
        fn=generate_scene,
        inputs=[
            input_image,
            scene_style,
            use_advanced,
            cfg_slider,
            steps_slider,
            denoise_slider
        ],
        outputs=[output_scene]
    )
    
    # Update rate limit status on load and after each generation
    demo.load(fn=get_rate_limit_status, inputs=None, outputs=[rate_status])
    btn_headshot.click(fn=get_rate_limit_status, inputs=None, outputs=[rate_status])
    btn_scene.click(fn=get_rate_limit_status, inputs=None, outputs=[rate_status])
    
    gr.Markdown(
        """
        ---
        ### ✨ Features
        - 🎭 **Background Remover**: Clean background removal for any image
        - 💼 **Professional Headshots**: Studio-quality headshots with multiple styles
        - 🌍 **Scene Changer**: Place yourself in different environments
        - ⚡ **AI-Powered**: InstantID technology for identity preservation
        - 🔒 **Rate Limited**: Fair usage with daily limits        
        
        """
    )

    gr.Markdown(
        """
        ---
        ### ⚠️ Notice
        - This is a **personal experimental project** and has undergone **limited testing**. Please use with caution.
        - Although **InstantID** is designed to preserve identity, outputs may vary based on **prompts, model weights, and other influencing factors**. 
        
        
        © 2025 Vijay S. Chaudhari | Powered by Replicate, InstantID and HuggingFace 🚀
        """
    )

# Launch
if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch()