Spaces:
Build error
Build error
| import torch | |
| from PIL import Image | |
| import gradio as gr | |
| import os | |
| from torchvision import transforms | |
| # Download the BeautyGAN repository directly from GitHub | |
| !pip install git+https://github.com/Honlan/BeautyGAN.git | |
| # Import the necessary model and utils | |
| from model import BeautyGAN | |
| # Initialize and load the pre-trained model | |
| model = BeautyGAN() | |
| model.load_state_dict(torch.load("models/beautygan.pth")) | |
| model.eval() | |
| # Define image transformation (to match input requirements of the model) | |
| transform = transforms.Compose([ | |
| transforms.Resize((256, 256)), # Resize image to fit the model's expected input size | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), # Normalize to [-1, 1] | |
| ]) | |
| # Inference function for BeautyGAN | |
| def beautygan_inference(image): | |
| # Apply the transformation to the image | |
| image_tensor = transform(image).unsqueeze(0) # Add batch dimension | |
| # Run the image through the BeautyGAN model | |
| with torch.no_grad(): | |
| output_image_tensor = model(image_tensor) # Get the output from the model | |
| # Convert the output tensor back to a PIL image | |
| output_image = output_image_tensor.squeeze(0).clamp(0, 1).permute(1, 2, 0).numpy() | |
| output_image = (output_image * 255).astype("uint8") | |
| output_image = Image.fromarray(output_image) | |
| return output_image | |
| # Create the Gradio interface | |
| iface = gr.Interface( | |
| fn=beautygan_inference, | |
| inputs=gr.Image(type="pil"), # Image input in PIL format | |
| outputs=gr.Image(type="pil"), # Image output in PIL format | |
| live=True # Enable live updates | |
| ) | |
| # Launch the Gradio interface | |
| iface.launch() | |