File size: 8,303 Bytes
8935f2f 1a8a785 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 |
---
# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
# Doc / guide: https://huggingface.co/docs/hub/model-cards
{}
---
# 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"
|