app.py

"""
🎨 Object Replacer AI - Versão Funcional
Substituição inteligente de objetos usando IA
"""

import gradio as gr
from PIL import Image, ImageDraw, ImageFilter, ImageFont
import numpy as np
import torch
from diffusers import StableDiffusionInpaintPipeline
from scipy.ndimage import binary_dilation
import tempfile
import time

print("🚀 Iniciando Object Replacer AI...")

# Configurações
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_SIZE = 768

print(f"💻 Dispositivo: {DEVICE}")

# Modelo global
inpaint_pipe = None

def log_msg(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}")

def load_inpaint_model():
    """Carrega modelo de inpainting"""
    global inpaint_pipe
    if inpaint_pipe is None:
        log_msg("📦 Carregando Stable Diffusion Inpainting...")
        try:
            inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                "stabilityai/stable-diffusion-2-inpainting",
                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
                safety_checker=None,
                requires_safety_checker=False
            ).to(DEVICE)
            
            if DEVICE == "cuda":
                inpaint_pipe.enable_attention_slicing()
                inpaint_pipe.enable_vae_slicing()
            
            log_msg("✅ Modelo carregado!")
        except Exception as e:
            log_msg(f"⚠️ Erro: {e}, tentando modelo alternativo...")
            inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                "runwayml/stable-diffusion-inpainting",
                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
                safety_checker=None
            ).to(DEVICE)
    return inpaint_pipe

def resize_image(image, max_size=MAX_SIZE):
    """Redimensiona mantendo proporção"""
    if max(image.size) <= max_size:
        return image
    ratio = max_size / max(image.size)
    new_size = (int(image.width * ratio), int(image.height * ratio))
    return image.resize(new_size, Image.Resampling.LANCZOS)

def create_circular_mask(image_size, center_x, center_y, radius):
    """Cria máscara circular ao redor do ponto clicado"""
    mask = Image.new('L', image_size, 0)
    draw = ImageDraw.Draw(mask)
    
    # Desenhar círculo
    draw.ellipse(
        [center_x - radius, center_y - radius, 
         center_x + radius, center_y + radius],
        fill=255
    )
    
    # Suavizar bordas
    mask = mask.filter(ImageFilter.GaussianBlur(radius=15))
    
    return mask

def create_mask_from_brush(image_size, brush_strokes):
    """Cria máscara a partir de traços de pincel"""
    mask = Image.new('L', image_size, 0)
    draw = ImageDraw.Draw(mask)
    
    if brush_strokes and len(brush_strokes) > 0:
        for stroke in brush_strokes:
            if len(stroke) >= 2:
                # Desenhar linha grossa
                draw.line(stroke, fill=255, width=30)
    
    # Suavizar
    mask = mask.filter(ImageFilter.GaussianBlur(radius=10))
    
    return mask

def smart_inpaint(image, mask, prompt, negative_prompt, steps, guidance):
    """Substitui objeto usando inpainting"""
    try:
        log_msg(f"🎨 Gerando: '{prompt}'")
        
        # Carregar modelo
        pipe = load_inpaint_model()
        
        # Ajustar tamanho para múltiplo de 8
        w, h = image.size
        new_w = (w // 8) * 8
        new_h = (h // 8) * 8
        
        img_resized = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
        mask_resized = mask.resize((new_w, new_h), Image.Resampling.LANCZOS)
        
        # Gerar
        start = time.time()
        
        with torch.no_grad():
            result = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                image=img_resized,
                mask_image=mask_resized,
                num_inference_steps=steps,
                guidance_scale=guidance,
                strength=0.99
            )
        
        output = result.images[0]
        
        # Voltar ao tamanho original
        output = output.resize((w, h), Image.Resampling.LANCZOS)
        
        elapsed = time.time() - start
        log_msg(f"✅ Concluído em {elapsed:.1f}s")
        
        return output, f"✅ Objeto gerado em {elapsed:.1f}s!"
        
    except Exception as e:
        log_msg(f"❌ Erro: {e}")
        return None, f"❌ Erro: {str(e)}"

def remove_object_smart(image, mask):
    """Remove objeto preenchendo com contexto"""
    try:
        log_msg("🗑️ Removendo objeto...")
        
        # Usar inpainting com prompt genérico
        prompt = "smooth background, natural surface, seamless, no objects"
        negative_prompt = "object, item, text, watermark, person, animal"
        
        result, msg = smart_inpaint(image, mask, prompt, negative_prompt, 40, 7.0)
        
        if result:
            return result, "✅ Objeto removido!"
        return None, msg
        
    except Exception as e:
        return None, f"❌ Erro ao remover: {str(e)}"

# Interface Gradio
with gr.Blocks(title="🎨 Object Replacer AI") as demo:
    
    gr.HTML("""
    <div style="text-align: center; padding: 25px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                border-radius: 12px; color: white; margin-bottom: 25px; box-shadow: 0 4px 15px rgba(0,0,0,0.2);">
        <h1 style="margin: 0; font-size: 2.5em; font-weight: 700;">✨ Object Replacer AI</h1>
        <p style="margin: 10px 0; font-size: 1.1em; opacity: 0.95;">Substitua ou remova objetos com Inteligência Artificial</p>
        <small style="opacity: 0.85;">Powered by Stable Diffusion Inpainting</small>
    </div>
    """)
    
    # Estado para armazenar dados
    mask_state = gr.State(None)
    original_image_state = gr.State(None)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 1. Carregue sua Imagem")
            input_image = gr.Image(
                label="",
                type="pil",
                interactive=True,
                height=350
            )
            
            gr.Markdown("### 🖌️ 2. Marque o Objeto")
            gr.Markdown("**Use o pincel para pintar sobre o objeto que deseja substituir/remover**")
            
            brush_image = gr.ImageMask(
                label="Pinte sobre o objeto",
                type="pil",
                height=350
            )
            
            gr.Markdown("**Dica:** Use o slider abaixo para ajustar o tamanho do traço")
            brush_size = gr.Slider(
                10, 50, 20, step=5,
                label="Tamanho do Pincel",
                info="Ajuste o tamanho do pincel"
            )
            
            clear_mask_btn = gr.Button("🔄 Limpar Máscara", size="sm")
        
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ 3. Configure a Ação")
            
            action_radio = gr.Radio(
                choices=[
                    ("🔄 Substituir por novo objeto", "replace"),
                    ("🗑️ Remover objeto", "remove")
                ],
                value="replace",
                label="O que fazer?",
                interactive=True
            )
            
            replacement_prompt = gr.Textbox(
                label="📝 Descreva o novo objeto (em inglês para melhor resultado)",
                placeholder="Ex: a red sports car / a beautiful flower / a modern laptop",
                lines=3,
                value="a beautiful red rose flower, photorealistic, high quality"
            )
            
            negative_prompt = gr.Textbox(
                label="🚫 O que evitar (opcional)",
                placeholder="Ex: ugly, blurry, low quality",
                value="ugly, blurry, low quality, distorted, deformed",
                lines=2
            )
            
            with gr.Accordion("🎛️ Configurações Avançadas", open=False):
                steps_slider = gr.Slider(
                    10, 50, 25, step=5,
                    label="Steps (Qualidade)",
                    info="Mais steps = melhor qualidade (mas mais lento)"
                )
                
                guidance_slider = gr.Slider(
                    1.0, 15.0, 7.5, step=0.5,
                    label="Guidance Scale",
                    info="7-9 recomendado. Maior = segue mais o prompt"
                )
            
            process_btn = gr.Button(
                "✨ PROCESSAR IMAGEM",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("### 🖼️ 4. Resultado")
            output_image = gr.Image(
                label="",
                type="pil",
                interactive=False,
                height=350
            )
            
            status_output = gr.Markdown("**Status:** Aguardando imagem...")
            
            with gr.Row():
                download_btn = gr.Button("📥 Baixar Resultado", size="lg")
                reset_btn = gr.Button("🔄 Resetar Tudo", variant="secondary", size="lg")
            
            download_file = gr.File(label="Download", visible=False)
    
    # Exemplos
    with gr.Accordion("💡 Exemplos de Prompts", open=False):
        gr.Markdown("""
        ### 🎯 Exemplos de prompts efetivos:
        
        **Para objetos realistas:**
        - `a red sports car, photorealistic, detailed, 4k`
        - `a beautiful yellow sunflower, vibrant colors, professional photo`
        - `a modern silver laptop, sleek design, high quality`
        - `a cute golden retriever puppy, fluffy, professional portrait`
        
        **Para estilo artístico:**
        - `a magical glowing crystal, fantasy art, ethereal lighting`
        - `a steampunk mechanical device, intricate gears, brass and copper`
        - `a neon cyberpunk sign, glowing letters, futuristic`
        
        **Dicas:**
        - ✅ Use inglês para melhores resultados
        - ✅ Seja específico: cores, estilo, qualidade
        - ✅ Adicione: "photorealistic, high quality, detailed"
        - ✅ Para remoção: deixe em branco ou use "smooth background"
        """)
    
    # Guia de uso
    gr.Markdown("""
    ---
    ### 📖 Como usar:
    
    1. **Carregue a imagem** - Clique ou arraste para a área de upload
    2. **Marque o objeto** - Use o pincel para pintar sobre o que quer substituir/remover
    3. **Escolha a ação**:
       - **Substituir**: Descreva o novo objeto em inglês
       - **Remover**: O objeto será removido e preenchido inteligentemente
    4. **Ajuste configurações** (opcional) - Steps e Guidance Scale
    5. **Clique em Processar** - Aguarde a IA gerar o resultado
    6. **Baixe o resultado** - Se gostar, clique em Baixar
    
    **⚠️ Importante:**
    - Primeira execução pode demorar (carregando modelo ~2GB)
    - Marque bem a área do objeto com o pincel
    - Use prompts em inglês para melhor qualidade
    - GPU acelera significativamente o processo
    """)
    
    # Funções de callback
    
    def update_brush_image(image):
        """Atualiza a imagem no brush quando nova imagem é carregada"""
        if image is None:
            return None, None, "❌ Carregue uma imagem primeiro"
        
        img = resize_image(image, MAX_SIZE)
        return img, img, "✅ Imagem carregada! Use o pincel para marcar o objeto"
    
    input_image.change(
        fn=update_brush_image,
        inputs=[input_image],
        outputs=[brush_image, original_image_state, status_output]
    )
    
    def toggle_prompt_visibility(action):
        """Mostra/esconde prompt baseado na ação"""
        return gr.update(visible=(action == "replace"))
    
    action_radio.change(
        fn=toggle_prompt_visibility,
        inputs=[action_radio],
        outputs=[replacement_prompt]
    )
    
    def process_image(original_img, brush_data, action, prompt, neg_prompt, steps, guidance):
        """Processa a imagem com inpainting"""
        try:
            log_msg("=== INÍCIO DO PROCESSAMENTO ===")
            
            if original_img is None:
                return None, "❌ Carregue uma imagem primeiro!"
            
            if brush_data is None:
                return None, "❌ Use o pincel para marcar o objeto!"
            
            # ImageMask do Gradio retorna a imagem com composição da máscara
            # Precisamos extrair apenas a máscara das áreas pintadas
            img = original_img
            
            # Verificar o formato do brush_data
            log_msg(f"Tipo brush_data: {type(brush_data)}")
            
            if isinstance(brush_data, dict):
                log_msg(f"Keys: {list(brush_data.keys())}")
                
                # Obter imagem composta (background + camadas)
                composite = brush_data.get('composite', None)
                background = brush_data.get('background', None) 
                layers = brush_data.get('layers', None)
                
                log_msg(f"Composite: {composite is not None}")
                log_msg(f"Background: {background is not None}")
                log_msg(f"Layers: {layers is not None}")
                
                if composite is not None:
                    # Usar a imagem composta
                    composite_img = composite if isinstance(composite, Image.Image) else Image.fromarray(composite)
                    
                    # A máscara é a diferença entre composite e background
                    if background is not None:
                        bg_img = background if isinstance(background, Image.Image) else Image.fromarray(background)
                        
                        # Converter para arrays numpy
                        composite_arr = np.array(composite_img.convert('RGB'))
                        bg_arr = np.array(bg_img.convert('RGB'))
                        
                        # Diferença absoluta
                        diff = np.abs(composite_arr.astype(int) - bg_arr.astype(int))
                        diff_sum = np.sum(diff, axis=2)
                        
                        # Criar máscara onde há diferença
                        mask_arr = (diff_sum > 30).astype(np.uint8) * 255
                        mask = Image.fromarray(mask_arr, mode='L')
                        
                        log_msg(f"Máscara criada por diferença: {mask.size}")
                    else:
                        # Se não tem background, tentar detectar as áreas pintadas
                        # Assumir que as áreas pintadas são escuras/pretas
                        composite_arr = np.array(composite_img.convert('L'))
                        mask_arr = (composite_arr < 128).astype(np.uint8) * 255
                        mask = Image.fromarray(mask_arr, mode='L')
                        
                        log_msg(f"Máscara criada por threshold: {mask.size}")
                
                elif layers is not None and len(layers) > 0:
                    # Processar camadas
                    log_msg(f"Processando {len(layers)} camadas")
                    
                    # Começar com máscara vazia
                    mask = Image.new('L', img.size, 0)
                    
                    for i, layer in enumerate(layers):
                        if layer is None:
                            continue
                            
                        log_msg(f"Layer {i}: {type(layer)}")
                        
                        if isinstance(layer, np.ndarray):
                            # Se for array, converter
                            if layer.ndim == 3:
                                # Se RGB/RGBA, pegar canal alpha ou converter para cinza
                                if layer.shape[2] == 4:
                                    layer_mask = Image.fromarray(layer[:, :, 3], mode='L')
                                else:
                                    layer_img = Image.fromarray(layer)
                                    layer_mask = layer_img.convert('L')
                            else:
                                layer_mask = Image.fromarray(layer, mode='L')
                        elif isinstance(layer, Image.Image):
                            layer_mask = layer.convert('L')
                        else:
                            continue
                        
                        # Combinar com máscara existente
                        mask_arr = np.array(mask)
                        layer_arr = np.array(layer_mask)
                        combined = np.maximum(mask_arr, layer_arr)
                        mask = Image.fromarray(combined, mode='L')
                
                else:
                    return None, "❌ Não foi possível extrair a máscara. Pinte sobre o objeto com o pincel!"
            
            elif isinstance(brush_data, Image.Image):
                # Se for uma imagem direta
                # A máscara é onde está pintado (normalmente em preto)
                brush_arr = np.array(brush_data.convert('L'))
                mask_arr = (brush_arr < 128).astype(np.uint8) * 255
                mask = Image.fromarray(mask_arr, mode='L')
                log_msg(f"Máscara de Image direta: {mask.size}")
            
            else:
                return None, "❌ Formato de máscara não reconhecido!"
            
            # Verificar se a máscara tem conteúdo
            if mask is None:
                return None, "❌ Não foi possível criar a máscara. Pinte sobre o objeto!"
            
            mask_array = np.array(mask)
            white_pixels = np.sum(mask_array > 128)
            total_pixels = mask_array.size
            percentage = (white_pixels / total_pixels) * 100
            
            log_msg(f"Pixels brancos: {white_pixels} ({percentage:.2f}% da imagem)")
            
            if white_pixels < 100:
                return None, f"❌ Área marcada muito pequena ({white_pixels} pixels). Pinte mais sobre o objeto!"
            
            # Dilatar a máscara levemente para cobrir melhor
            mask_array = np.array(mask)
            kernel_size = 5
            kernel = np.ones((kernel_size, kernel_size), np.uint8)
            
            from scipy.ndimage import binary_dilation
            mask_dilated = binary_dilation(mask_array > 128, structure=kernel).astype(np.uint8) * 255
            mask = Image.fromarray(mask_dilated, mode='L')
            
            # Suavizar bordas
            mask = mask.filter(ImageFilter.GaussianBlur(radius=5))
            
            # Redimensionar
            img = resize_image(img, MAX_SIZE)
            mask = mask.resize(img.size, Image.Resampling.LANCZOS)
            
            log_msg(f"✅ Máscara válida! Processando - Ação: {action}")
            
            # Processar baseado na ação
            if action == "replace":
                if not prompt or len(prompt.strip()) < 3:
                    return None, "❌ Descreva o novo objeto no campo de texto!"
                
                result, status = smart_inpaint(img, mask, prompt, neg_prompt, steps, guidance)
            else:  # remove
                result, status = remove_object_smart(img, mask)
            
            return result, status
            
        except Exception as e:
            import traceback
            error = f"❌ Erro: {str(e)}\n\n{traceback.format_exc()}"
            log_msg(error)
            return None, error
    
    process_btn.click(
        fn=process_image,
        inputs=[original_image_state, brush_image, action_radio, 
                replacement_prompt, negative_prompt, steps_slider, guidance_slider],
        outputs=[output_image, status_output]
    )
    
    def save_result(image):
        """Salva resultado"""
        if image is None:
            return None
        
        temp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        image.save(temp.name, "PNG", optimize=True)
        return temp.name
    
    download_btn.click(
        fn=save_result,
        inputs=[output_image],
        outputs=[download_file]
    )
    
    def reset_all():
        """Reseta tudo"""
        return None, None, None, None, None, "🔄 Tudo resetado! Comece novamente"
    
    reset_btn.click(
        fn=reset_all,
        outputs=[input_image, brush_image, output_image, mask_state, 
                original_image_state, status_output]
    )
    
    def clear_mask(img):
        """Limpa apenas a máscara"""
        if img:
            return img
        return None
    
    clear_mask_btn.click(
        fn=clear_mask,
        inputs=[original_image_state],
        outputs=[brush_image]
    )
    
    gr.HTML("""
    <div style="text-align: center; margin-top: 30px; padding: 20px; color: #666; 
                border-top: 1px solid #e0e0e0;">
        <p style="margin: 5px 0;"><strong>✨ Object Replacer AI ✨</strong></p>
        <p style="font-size: 0.9em; margin: 5px 0;">
            🎨 Stable Diffusion Inpainting • 🖌️ Interface Intuitiva • ⚡ Resultados Profissionais
        </p>
        <p style="font-size: 0.8em; color: #999; margin-top: 10px;">
            Dica: Para melhores resultados, marque bem o objeto e use prompts descritivos em inglês
        </p>
    </div>
    """)

if __name__ == "__main__":
    log_msg(f"✅ Sistema pronto! Dispositivo: {DEVICE}")
    log_msg("⚡ Modelo será carregado na primeira geração")
    
    demo.launch(
        server_name="0.0.0.0",
        share=True,
        show_error=True,
        theme=gr.themes.Soft()
    )