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()