AI-generated Art Creator Python, AI

```python
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets, models
from PIL import Image
import os
import random
import numpy as np
import matplotlib.pyplot as plt

# Check for CUDA availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")


# -------------------- Utility Functions --------------------

def load_image(image_path, transform=None):
    """Loads an image and applies optional transformations."""
    try:
        image = Image.open(image_path).convert('RGB')  # Ensure RGB for consistency
        if transform:
            image = transform(image).unsqueeze(0)  # Add batch dimension
        return image.to(device)
    except FileNotFoundError:
        print(f"Error: Image file not found at {image_path}")
        return None
    except Exception as e:
        print(f"Error loading image: {e}")
        return None

def save_image(tensor, filename):
    """Saves a tensor as an image file."""
    image = tensor.cpu().clone()
    image = image.squeeze(0)  # Remove batch dimension
    image = unloader(image)
    image.save(filename)



# -------------------- Style Transfer Model --------------------

class StyleTransferModel(nn.Module):
    def __init__(self, style_image_path, content_layers=None, style_layers=None):
        super(StyleTransferModel, self).__init__()

        self.vgg = models.vgg19(pretrained=True).features.to(device).eval()
        for param in self.vgg.parameters():
            param.requires_grad_(False)

        self.content_layers = content_layers or ['conv_4']
        self.style_layers = style_layers or ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']

        self.style_weights = {'conv_1': 1.0, 'conv_2': 0.8, 'conv_3': 0.5, 'conv_4': 0.3, 'conv_5': 0.1}  # Adjust as needed

        self.style_image = load_image(style_image_path, transform=transform)
        if self.style_image is None:
            raise ValueError("Failed to load style image. Check path and image format.")

        self.style_features = self.get_features(self.style_image, self.style_layers)

    def get_features(self, image, layers):
        """Extracts features from specified layers of the VGG network."""
        features = {}
        x = image
        for name, layer in self.vgg._modules.items():
            x = layer(x)
            if name in layer_mappings:
                mapped_name = layer_mappings[name]  # Use the dictionary
                if mapped_name in layers:
                    features[mapped_name] = x
        return features

    def gram_matrix(self, tensor):
        """Calculates the Gram matrix of a given tensor."""
        batch_size, depth, height, width = tensor.size()
        tensor = tensor.view(batch_size * depth, height * width)
        gram = torch.mm(tensor, tensor.t())
        return gram.div(batch_size * depth * height * width)

    def content_loss(self, content_features, generated_features, layer):
        """Calculates the content loss."""
        return torch.mean((generated_features[layer] - content_features[layer]) ** 2)

    def style_loss(self, generated_features, layer):
        """Calculates the style loss."""
        style_gram = self.gram_matrix(self.style_features[layer])
        generated_gram = self.gram_matrix(generated_features[layer])
        return self.style_weights[layer] * torch.mean((generated_gram - style_gram) ** 2)


    def forward(self, generated_image, content_image):
        """Performs the style transfer and calculates the total loss."""

        content_features = self.get_features(content_image, self.content_layers)
        generated_features = self.get_features(generated_image, self.content_layers + self.style_layers)

        total_content_loss = 0
        for layer in self.content_layers:
            total_content_loss += self.content_loss(content_features, generated_features, layer)

        total_style_loss = 0
        for layer in self.style_layers:
            total_style_loss += self.style_loss(generated_features, layer)

        return total_content_loss, total_style_loss


# -------------------- Training Loop --------------------

def train_style_transfer(content_image_path, style_image_path, output_image_path,
                         num_epochs=2000, content_weight=1, style_weight=1e5):
    """Trains the style transfer model and generates the stylized image."""

    content_image = load_image(content_image_path, transform=transform)
    if content_image is None:
        print("Failed to load content image. Check path and image format.")
        return

    # Initialize the generated image (start from the content image)
    generated_image = content_image.clone().requires_grad_(True)
    optimizer = optim.Adam([generated_image], lr=0.01)

    # Initialize the StyleTransferModel here, after loading the content image
    model = StyleTransferModel(style_image_path).to(device)


    for epoch in range(num_epochs + 1): # Include epoch 0 for initial save
        content_loss, style_loss = model(generated_image, content_image)
        total_loss = content_weight * content_loss + style_weight * style_loss

        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()

        if epoch % 100 == 0:
            print(f"Epoch {epoch}: Total Loss = {total_loss.item():.4f}, Content Loss = {content_loss.item():.4f}, Style Loss = {style_loss.item():.4f}")
            save_image(generated_image, f"{output_image_path.replace('.jpg', '')}_epoch_{epoch}.jpg") # Save at each checkpoint

    # Save the final image
    save_image(generated_image, output_image_path)
    print(f"Final image saved to {output_image_path}")


# -------------------- Image Preprocessing --------------------

image_size = 256  # You can adjust the image size here
transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

unloader = transforms.ToPILImage()  # For converting back to PIL image

# Layer mappings for VGG
layer_mappings = {
    '0': 'conv_1',
    '2': 'conv_2',
    '5': 'conv_3',
    '7': 'conv_4',
    '10': 'conv_5',
    '12': 'conv_6',
    '14': 'conv_7',
    '16': 'conv_8',
    '19': 'conv_9',
    '21': 'conv_10',
    '23': 'conv_11',
    '25': 'conv_12',
    '28': 'conv_13',
    '30': 'conv_14',
    '32': 'conv_15',
    '34': 'conv_16'
}

# -------------------- Main Execution --------------------

if __name__ == "__main__":
    # Example Usage:
    content_image_path = "content.jpg"  # Replace with your content image path
    style_image_path = "style.jpg"      # Replace with your style image path
    output_image_path = "output.jpg"    # Replace with your desired output path

    # Create dummy images if they don't exist
    if not os.path.exists(content_image_path):
        dummy_image = Image.new('RGB', (256, 256), color='white')
        dummy_image.save(content_image_path)
        print(f"Created dummy content image at {content_image_path}")
    if not os.path.exists(style_image_path):
        dummy_image = Image.new('RGB', (256, 256), color='blue')
        dummy_image.save(style_image_path)
        print(f"Created dummy style image at {style_image_path}")


    try:
        train_style_transfer(content_image_path, style_image_path, output_image_path, num_epochs=1000)  # You can adjust the number of epochs
    except Exception as e:
        print(f"An error occurred during training: {e}")
```