metadata
{}
Model Card for Model ID
copy/paste/save as pix2pixinference.py
import argparse
import torch
from torchvision.transforms import Compose, Resize, ToTensor, Normalize, ToPILImage
from torchvision.utils import save_image
from PIL import Image
import os
import io
from huggingface_hub import hf_hub_download
import sys
import matplotlib.pyplot as plt
# Import the model architecture - assuming it's locally available
# If not, we'll need to define it here
try:
from modeling_pix2pix import GeneratorUNet
except ImportError:
print("Couldn't import model architecture, defining it here...")
# Define the UNet architecture as it appears in the original code
import torch.nn as nn
import torch.nn.functional as F
def weights_init_normal(m):
classname = m.__class__.__name__
if classname.find("Conv") != -1:
torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
elif classname.find("BatchNorm2d") != -1:
torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
torch.nn.init.constant_(m.bias.data, 0.0)
class UNetDown(nn.Module):
def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
super(UNetDown, self).__init__()
layers = [nn.Conv2d(in_size, out_size, 4, 2, 1, bias=False)]
if normalize:
layers.append(nn.InstanceNorm2d(out_size))
layers.append(nn.LeakyReLU(0.2))
if dropout:
layers.append(nn.Dropout(dropout))
self.model = nn.Sequential(*layers)
def forward(self, x):
return self.model(x)
class UNetUp(nn.Module):
def __init__(self, in_size, out_size, dropout=0.0):
super(UNetUp, self).__init__()
layers = [
nn.ConvTranspose2d(in_size, out_size, 4, 2, 1, bias=False),
nn.InstanceNorm2d(out_size),
nn.ReLU(inplace=True),
]
if dropout:
layers.append(nn.Dropout(dropout))
self.model = nn.Sequential(*layers)
def forward(self, x, skip_input):
x = self.model(x)
x = torch.cat((x, skip_input), 1)
return x
class GeneratorUNet(nn.Module):
def __init__(self, in_channels=3, out_channels=3):
super(GeneratorUNet, self).__init__()
self.down1 = UNetDown(in_channels, 64, normalize=False)
self.down2 = UNetDown(64, 128)
self.down3 = UNetDown(128, 256)
self.down4 = UNetDown(256, 512, dropout=0.5)
self.down5 = UNetDown(512, 512, dropout=0.5)
self.down6 = UNetDown(512, 512, dropout=0.5)
self.down7 = UNetDown(512, 512, dropout=0.5)
self.down8 = UNetDown(512, 512, normalize=False, dropout=0.5)
self.up1 = UNetUp(512, 512, dropout=0.5)
self.up2 = UNetUp(1024, 512, dropout=0.5)
self.up3 = UNetUp(1024, 512, dropout=0.5)
self.up4 = UNetUp(1024, 512, dropout=0.5)
self.up5 = UNetUp(1024, 256)
self.up6 = UNetUp(512, 128)
self.up7 = UNetUp(256, 64)
self.final = nn.Sequential(
nn.ConvTranspose2d(128, out_channels, 4, 2, 1),
nn.Tanh(),
)
def forward(self, x):
# U-Net generator with skip connections from encoder to decoder
d1 = self.down1(x)
d2 = self.down2(d1)
d3 = self.down3(d2)
d4 = self.down4(d3)
d5 = self.down5(d4)
d6 = self.down6(d5)
d7 = self.down7(d6)
d8 = self.down8(d7)
u1 = self.up1(d8, d7)
u2 = self.up2(u1, d6)
u3 = self.up3(u2, d5)
u4 = self.up4(u3, d4)
u5 = self.up5(u4, d3)
u6 = self.up6(u5, d2)
u7 = self.up7(u6, d1)
return self.final(u7)
def parse_args():
parser = argparse.ArgumentParser(description="Generate images using Pix2Pix model from HuggingFace Hub")
parser.add_argument(
"--repo_id",
type=str,
required=True,
help="HuggingFace Hub repository ID (e.g., 'username/model_name')"
)
parser.add_argument(
"--model_file",
type=str,
default="model.pt",
help="Name of the model file in the repository"
)
parser.add_argument(
"--input_image",
type=str,
required=True,
help="Path to input image (night image to transform to day)"
)
parser.add_argument(
"--output_image",
type=str,
default="output.png",
help="Path to save the generated image"
)
parser.add_argument(
"--image_size",
type=int,
default=256,
help="Size of the input/output images"
)
parser.add_argument(
"--display",
action="store_true",
help="Display input and output images using matplotlib"
)
parser.add_argument(
"--token",
type=str,
default=None,
help="HuggingFace token for accessing private repositories"
)
return parser.parse_args()
def main():
args = parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Set up image transformations
transform_input = Compose([
Resize((args.image_size, args.image_size)),
ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
# Initialize model
print("Initializing model...")
generator = GeneratorUNet()
generator.to(device)
# Download model from Hugging Face Hub
print(f"Downloading model from {args.repo_id}...")
try:
model_path = hf_hub_download(
repo_id=args.repo_id,
filename=args.model_file,
token=args.token
)
print(f"Model downloaded to {model_path}")
except Exception as e:
print(f"Error downloading model: {e}")
sys.exit(1)
# Load model weights
try:
generator.load_state_dict(torch.load(model_path, map_location=device))
generator.eval()
print("Model loaded successfully")
except Exception as e:
print(f"Error loading model weights: {e}")
sys.exit(1)
# Load and preprocess input image
try:
image = Image.open(args.input_image).convert("RGB")
original_image = image.copy()
input_tensor = transform_input(image).unsqueeze(0).to(device)
print(f"Input image loaded: {args.input_image}")
except Exception as e:
print(f"Error loading input image: {e}")
sys.exit(1)
# Generate output image
print("Generating image...")
with torch.no_grad():
fake_B = generator(input_tensor)
# Save the output image
try:
# Denormalize and convert back to image
output_image = fake_B.cpu()
save_image(output_image, args.output_image, normalize=True)
print(f"Output image saved to {args.output_image}")
# Create a PIL image for display if needed
to_pil = ToPILImage()
output_pil = to_pil(output_image.squeeze(0) * 0.5 + 0.5)
except Exception as e:
print(f"Error saving output image: {e}")
sys.exit(1)
# Display images if requested
if args.display:
try:
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.title("Input Image (Night)")
plt.imshow(original_image)
plt.axis("off")
plt.subplot(1, 2, 2)
plt.title("Generated Image (Day)")
plt.imshow(output_pil)
plt.axis("off")
plt.tight_layout()
plt.show()
except Exception as e:
print(f"Error displaying images: {e}")
if __name__ == "__main__":
main()
python pix2pixinference.py --repo_id "uisikdag/gan-pix2pix-night2day" --input_image "night_image.jpg" --output_image "day_image.png"