import gradio as gr import spaces import tensorflow as tf import numpy as np from PIL import Image import logging import time from tqdm import tqdm # Initialize logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger("style_transfer_app") # Set TensorFlow threading options tf.config.threading.set_inter_op_parallelism_threads(8) tf.config.threading.set_intra_op_parallelism_threads(8) def load_img(image): """Load and preprocess image for style transfer""" max_dim = 512 # Convert PIL Image to tensor img = tf.convert_to_tensor(np.array(image)) img = tf.image.convert_image_dtype(img, tf.float32) shape = tf.cast(tf.shape(img)[:-1], tf.float32) long_dim = max(shape) scale = max_dim / long_dim new_shape = tf.cast(shape * scale, tf.int32) img = tf.image.resize(img, new_shape) img = img[tf.newaxis, :] return img def vgg_layers(layer_names): """Create VGG model with specified layers""" vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet') vgg.trainable = False outputs = [vgg.get_layer(name).output for name in layer_names] model = tf.keras.Model([vgg.input], outputs) return model def gram_matrix(input_tensor): """Calculate Gram matrix""" result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor) input_shape = tf.shape(input_tensor) num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32) return result / num_locations class StyleContentModel(tf.keras.models.Model): def __init__(self, style_layers, content_layers): super(StyleContentModel, self).__init__() self.vgg = vgg_layers(style_layers + content_layers) self.style_layers = style_layers self.content_layers = content_layers self.num_style_layers = len(style_layers) self.vgg.trainable = False def call(self, inputs): inputs = inputs * 255.0 preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs) outputs = self.vgg(preprocessed_input) style_outputs, content_outputs = (outputs[:self.num_style_layers], outputs[self.num_style_layers:]) style_outputs = [gram_matrix(style_output) for style_output in style_outputs] content_dict = {content_name: value for content_name, value in zip(self.content_layers, content_outputs)} style_dict = {style_name: value for style_name, value in zip(self.style_layers, style_outputs)} return {'content': content_dict, 'style': style_dict} def clip_0_1(image): """Clip tensor values between 0 and 1""" return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0) def style_content_loss(outputs, style_targets, content_targets, style_weight, content_weight): """Calculate style and content loss""" style_outputs = outputs['style'] content_outputs = outputs['content'] style_loss = tf.add_n([tf.reduce_mean((style_outputs[name]-style_targets[name])**2) for name in style_outputs.keys()]) style_loss *= style_weight / len(style_outputs) content_loss = tf.add_n([tf.reduce_mean((content_outputs[name]-content_targets[name])**2) for name in content_outputs.keys()]) content_loss *= content_weight / len(content_outputs) loss = style_loss + content_loss return loss @tf.function def train_step(image, extractor, style_targets, content_targets, opt, style_weight, content_weight, total_variation_weight): """Perform one training step""" with tf.GradientTape() as tape: outputs = extractor(image) loss = style_content_loss(outputs, style_targets, content_targets, style_weight, content_weight) loss += total_variation_weight * tf.image.total_variation(image) grad = tape.gradient(loss, image) opt.apply_gradients([(grad, image)]) image.assign(clip_0_1(image)) return loss def tensor_to_image(tensor): """Convert tensor to PIL Image""" tensor = tensor * 255 tensor = np.array(tensor, dtype=np.uint8) if np.ndim(tensor) > 3: tensor = tensor[0] return Image.fromarray(tensor) # Initialize the style-content model style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1'] content_layers = ['block5_conv2'] extractor = StyleContentModel(style_layers, content_layers) @spaces.GPU(duration=120) # Style transfer typically needs more than 60s def style_transfer_fn(content_image, style_image, progress=gr.Progress(track_tqdm=True)): """Main style transfer function for Gradio interface""" try: # Preprocess images content_img = load_img(content_image) style_img = load_img(style_image) # Extract style and content features style_targets = extractor(style_img)['style'] content_targets = extractor(content_img)['content'] image = tf.Variable(content_img) # Set optimization parameters opt = tf.keras.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1) style_weight = 1e-2 content_weight = 1e4 total_variation_weight = 30 epochs = 10 steps_per_epoch = 100 start_time = time.time() # Training loop for n in tqdm(range(epochs), desc="Epochs"): for m in tqdm(range(steps_per_epoch), desc="Steps", leave=False): loss = train_step(image, extractor, style_targets, content_targets, opt, style_weight, content_weight, total_variation_weight) # Convert result to image result_image = tensor_to_image(image) return result_image except Exception as e: logger.error(f"Error during style transfer: {e}") raise gr.Error("An error occurred during style transfer.") # Create Gradio interface iface = gr.Interface( fn=style_transfer_fn, inputs=[ gr.Image(label="Content Image", type="pil"), gr.Image(label="Style Image", type="pil") ], outputs=gr.Image(label="Stylized Image"), title="Neural Style Transfer - Ty Chermsirivatana", description="Upload a content image and a style image to create a stylized image in context.", ) # Launch the interface if __name__ == "__main__": iface.launch()