# ==========================================
# PATCH DE COMPATIBILIDADE (TRANSFORMERS 5.X + RMBG 1.4)
# Deve ficar no TOPO do arquivo, antes de tudo!
# ==========================================
import torch
_original_getattr = torch.nn.Module.__getattr__

def _patched_getattr(self, name):
    if name == "all_tied_weights_keys":
        return {}
    return _original_getattr(self, name)

torch.nn.Module.__getattr__ = _patched_getattr
# ==========================================
import gradio as gr
from transformers import pipeline
from PIL import Image
import spaces

# 1. Defina o dispositivo padrão (GPU se disponível, senão CPU)
# Nota: No ZeroGPU, isso inicializa na CPU, o spaces.GPU gerencia a mudança.
device = 0 if torch.cuda.is_available() else -1

print(f"Inicializando pipeline no dispositivo: {'GPU' if device == 0 else 'CPU'}")


# Função para remover o background da imagem
@spaces.GPU
def remove_background(image):
    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=device)

    # Obter a máscara da imagem
    pillow_mask = pipe(image, return_mask=True)

    # Aplicar máscara na imagem original
    pillow_image = pipe(image)

    return pillow_image

# Criar uma interface Gradio
app = gr.Interface(
    fn=remove_background,
    inputs=gr.components.Image(type="pil"),
    outputs=gr.components.Image(type="pil", format="png"),  # Especificar saída como PNG
    title="Remoção de Background de Imagens",
    description="Envie uma imagem e veja o background sendo removido automaticamente. A imagem resultante será no formato PNG."
)

if __name__ == '__main__':
    app.launch()