ssm1986/codeformer-api

1import gradio as gr2import torch3from basicsr.archs.codeformer_arch import CodeFormer4from torchvision.transforms.functional import to_tensor5from PIL import Image6 7# Load CodeFormer model8device = 'cuda' if torch.cuda.is_available() else 'cpu'9net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9).to(device)10ckpt_path = 'CodeFormer/codeformer.pth'  # Path to the model checkpoint11 12net.load_state_dict(torch.load(ckpt_path, map_location=device)['params_ema'])13net.eval()14 15def enhance_image(input_image):16    img_tensor = to_tensor(input_image).unsqueeze(0).to(device)17    with torch.no_grad():18        output = net(img_tensor, w=0.7, adain=True)[0].cpu()19    enhanced_image = Image.fromarray((output.clamp(0, 1).numpy().transpose(1, 2, 0) * 255).astype('uint8'))20    return enhanced_image21 22iface = gr.Interface(fn=enhance_image, inputs=gr.Image(type="pil"), outputs=gr.Image(), title="CodeFormer Enhancer")23iface.launch()24