init commit

app.py

@@ -0,0 +1,133 @@
+import streamlit as st
+import torch
+import numpy as np
+from PIL import Image
+from model import CycleGAN, get_val_transform, de_normalize
+
+# Configure page
+st.set_page_config(
+    page_title="CycleGAN Image Converter",
+    page_icon="🎨",
+    layout="wide"
+)
+
+# Get the best available device
+@st.cache_resource
+def get_device():
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        st.sidebar.success("Using GPU 🚀")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+        st.sidebar.success("Using Apple Silicon 🍎")
+    else:
+        device = torch.device("cpu")
+        st.sidebar.info("Using CPU 💻")
+    return device
+
+# Add custom CSS
+st.markdown("""
+    <style>
+    .stApp {
+        max-width: 1200px;
+        margin: 0 auto;
+    }
+    .main {
+        padding: 2rem;
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+# Title and description
+st.title("CycleGAN Image Converter 🎨")
+st.markdown("""
+    Transform images between different domains using CycleGAN.
+    Upload an image and see it converted in real-time!
+    
+    *Note: Images will be resized to 256x256 pixels during conversion.*
+""")
+
+# Available models and their configurations
+MODELS = [
+    {
+        "name": "Cezanne ↔ Photo",
+        "id": "cezanne2photo",
+        "model_path": "waleko/cyclegan",
+        "description": "Convert between Cezanne's painting style and photographs"
+    }
+]
+
+# Sidebar controls
+with st.sidebar:
+    st.header("Settings")
+    
+    # Model selection
+    selected_model = st.selectbox(
+        "Conversion Type",
+        options=range(len(MODELS)),
+        format_func=lambda x: MODELS[x]["name"]
+    )
+    
+    # Direction selection
+    direction = st.radio(
+        "Conversion Direction",
+        options=["A → B", "B → A"],
+        help="A → B: Convert from domain A to B\nB → A: Convert from domain B to A"
+    )
+
+# Load model
+@st.cache_resource
+def load_model(model_path):
+    device = get_device()
+    model = CycleGAN.from_pretrained(model_path)
+    model = model.to(device)
+    model.eval()
+    return model
+
+# Process image
+def process_image(image, model, direction):
+    # Prepare transform
+    transform = get_val_transform(model, direction)
+    
+    # Convert PIL image to tensor
+    tensor = transform(np.array(image)).unsqueeze(0)
+    
+    # Move to appropriate device
+    tensor = tensor.to(next(model.parameters()).device)
+    
+    # Process
+    with torch.no_grad():
+        if direction == "A → B":
+            output = model.generator_ab(tensor)
+        else:
+            output = model.generator_ba(tensor)
+    
+    # Convert back to image
+    result = de_normalize(output[0], model, direction)
+    return result
+
+# Main interface
+col1, col2 = st.columns(2)
+
+with col1:
+    st.subheader("Input Image")
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    
+    if uploaded_file is not None:
+        input_image = Image.open(uploaded_file)
+        st.image(input_image, use_column_width=True)
+
+with col2:
+    st.subheader("Converted Image")
+    if uploaded_file is not None:
+        try:
+            # Load and process
+            model = load_model(MODELS[selected_model]["model_path"])
+            result = process_image(input_image, model, direction)
+            
+            # Display
+            st.image(result, use_column_width=True)
+        except Exception as e:
+            st.error(f"Error during conversion: {str(e)}")
+    else:
+        st.info("Upload an image to see the conversion result") 

model.py

@@ -0,0 +1,95 @@
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+import torchvision.transforms as tr
+import functools
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+
+        self.block = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+def generator(num_residual_blocks=9):
+    channels = 3
+    out_features = 64
+    model = [
+        nn.ReflectionPad2d(channels),
+        nn.Conv2d(channels, out_features, 7),
+        nn.InstanceNorm2d(out_features),
+        nn.ReLU(inplace=True),
+    ]
+    in_features = out_features
+
+    # Downsampling
+    for _ in range(2):
+        out_features *= 2
+        model += [
+            nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+            nn.InstanceNorm2d(out_features),
+            nn.ReLU(inplace=True),
+        ]
+        in_features = out_features
+
+    # Residual blocks
+    for _ in range(num_residual_blocks):
+        model += [ResidualBlock(out_features)]
+
+    # Upsampling
+    for _ in range(2):
+        out_features //= 2
+        model += [
+            nn.Upsample(scale_factor=2),
+            nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),
+            nn.InstanceNorm2d(out_features),
+            nn.ReLU(inplace=True),
+        ]
+        in_features = out_features
+
+    # Output layer
+    model += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]
+
+    return nn.Sequential(*model)
+
+class CycleGAN(nn.Module, PyTorchModelHubMixin, pipeline_tag="image-to-image"):
+    def __init__(self, channel_mean_a=None, channel_std_a=None, channel_mean_b=None, channel_std_b=None):
+        super(CycleGAN, self).__init__()
+        self.generator_ab = generator()
+        self.generator_ba = generator()
+        
+        # Store normalization parameters as non-trainable parameters
+        self.register_buffer('channel_mean_a', torch.tensor(channel_mean_a if channel_mean_a is not None else [0.5, 0.5, 0.5], dtype=torch.float32))
+        self.register_buffer('channel_std_a', torch.tensor(channel_std_a if channel_std_a is not None else [0.5, 0.5, 0.5], dtype=torch.float32))
+        self.register_buffer('channel_mean_b', torch.tensor(channel_mean_b if channel_mean_b is not None else [0.5, 0.5, 0.5], dtype=torch.float32))
+        self.register_buffer('channel_std_b', torch.tensor(channel_std_b if channel_std_b is not None else [0.5, 0.5, 0.5], dtype=torch.float32))
+
+def get_val_transform(model, direction="a_to_b", size=256):
+    mean = model.channel_mean_a if direction == "a_to_b" else model.channel_mean_b
+    std = model.channel_std_a if direction == "a_to_b" else model.channel_std_b
+    
+    return tr.Compose([
+        tr.ToPILImage(),
+        tr.Resize(size),
+        tr.CenterCrop(size),
+        tr.ToTensor(),
+        tr.Normalize(mean=mean.tolist(), std=std.tolist()),
+    ])
+
+def de_normalize(tensor, model, direction="a_to_b"):
+    img_tensor = tensor.cpu().detach().clone()
+    mean = model.channel_mean_a if direction == "a_to_b" else model.channel_mean_b
+    std = model.channel_std_a if direction == "a_to_b" else model.channel_std_b
+    
+    img_tensor = img_tensor * std[:, None, None] + mean[:, None, None]
+    return torch.clamp(img_tensor.permute(1, 2, 0), 0.0, 1.0) 

requirements.txt

@@ -0,0 +1,6 @@
+streamlit>=1.28.0
+torch>=2.0.0
+torchvision>=0.15.0
+Pillow>=10.0.0
+numpy>=1.24.0
+huggingface-hub>=0.19.0