{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "fbdfdd2b", "metadata": {}, "outputs": [], "source": [ "#configuring the system environment\n", "!conda create --name diffusion python=3.9\n", "!conda activate diffusion\n", "!pip install ipykernel\n", "!python -m ipykernel install --user --name diffusion --display-name \"Python (diffusion)\"\n", "!pip install -r sliders/requirements.txt\n", "!pip install tensorflow tensorflow_datasets\n", "!pip install torchmetrics\n", "!pip install lpips" ] }, { "cell_type": "code", "execution_count": null, "id": "c70205ee", "metadata": {}, "outputs": [], "source": [ "import os\n", "try:\n", " os.chdir('sliders') # Attempt to change the directory to 'sliders'\n", "except:\n", " pass # If the directory doesn't exist, continue without error\n", "import torch\n", "from PIL import Image\n", "import random\n", "import pandas as pd\n", "import matplotlib.pyplot as plt\n", "import gc\n", "from transformers import CLIPProcessor, CLIPModel # Import CLIP model and processor\n", "import diffusers\n", "from diffusers import DiffusionPipeline\n", "from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipeline, StableDiffusionXLPipelineOutput\n", "from trainscripts.textsliders.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV\n", "import torch.nn as nn\n", "from typing import Any, Callable, Dict, List, Optional, Tuple, Union\n", "import torchvision.transforms as transforms\n", "from torchvision import models, datasets\n", "import lpips # Import LPIPS (Learned Perceptual Image Patch Similarity) for perceptual similarity measurement\n", "\n", "def flush():\n", " \"\"\"Function to clear the CUDA memory cache and perform garbage collection.\"\"\"\n", " torch.cuda.empty_cache() # Clear the CUDA memory cache\n", " gc.collect() # Run garbage collection\n", "\n", "# Define the Discriminator network for the GAN\n", "class Discriminator(nn.Module):\n", " def __init__(self, input_size):\n", " super(Discriminator, self).__init__()\n", " # Define convolutional layers with LeakyReLU activation\n", " self.conv1 = nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1)\n", " self.conv2 = nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1)\n", " # Fully connected layer for final output\n", " self.fc1 = nn.Linear(128 * (input_size // 4) * (input_size // 4), 1)\n", "\n", " def forward(self, x):\n", " # Apply LeakyReLU activation to each layer\n", " x = nn.LeakyReLU(0.2)(self.conv1(x))\n", " x = nn.LeakyReLU(0.2)(self.conv2(x))\n", " x = x.view(x.size(0), -1) # Flatten the output\n", " x = torch.sigmoid(self.fc1(x)) # Apply sigmoid activation to the output\n", " return x\n", "\n", "# Define the Perceptual Loss network using a pre-trained VGG16 model\n", "class PerceptualLoss(nn.Module):\n", " def __init__(self):\n", " super(PerceptualLoss, self).__init__()\n", " # Load the first 16 layers of VGG16 for feature extraction\n", " vgg = models.vgg16(pretrained=True).features[:16]\n", " self.vgg_layers = vgg\n", " # Freeze the parameters of VGG16\n", " for param in self.vgg_layers.parameters():\n", " param.requires_grad = False\n", "\n", " def forward(self, x, y):\n", " # Compute the VGG16 feature representations of the input images\n", " x_vgg = self.vgg_layers(x)\n", " y_vgg = self.vgg_layers(y)\n", " # Calculate the mean squared error between the features of the two images\n", " loss = nn.functional.mse_loss(x_vgg, y_vgg)\n", " return loss\n", "\n", "# Load the CIFAR-10 dataset\n", "transform = transforms.Compose([\n", " transforms.Resize((1024, 1024)), # Resize images to 1024x1024 pixels\n", " transforms.ToTensor() # Convert images to tensor format\n", "])\n", "cifar10 = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\n", "cifar10_loader = torch.utils.data.DataLoader(cifar10, batch_size=64, shuffle=True)\n", "real_images = next(iter(cifar10_loader))[0]\n", "real_images = real_images.requires_grad_(True) # Ensure that real_images have gradient tracking enabled\n", "\n", "@torch.no_grad()\n", "def call(\n", " self,\n", " prompt: Union[str, List[str]] = None,\n", " prompt_2: Optional[Union[str, List[str]]] = None,\n", " height: Optional[int] = None,\n", " width: Optional[int] = None,\n", " num_inference_steps: int = 50,\n", " denoising_end: Optional[float] = None,\n", " guidance_scale: float = 5.0,\n", " negative_prompt: Optional[Union[str, List[str]]] = None,\n", " negative_prompt_2: Optional[Union[str, List[str]]] = None,\n", " num_images_per_prompt: Optional[int] = 1,\n", " eta: float = 0.0,\n", " generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n", " latents: Optional[torch.FloatTensor] = None,\n", " prompt_embeds: Optional[torch.FloatTensor] = None,\n", " negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n", " pooled_prompt_embeds: Optional[torch.FloatTensor] = None,\n", " negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,\n", " output_type: Optional[str] = \"pil\",\n", " return_dict: bool = True,\n", " callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n", " callback_steps: int = 1,\n", " cross_attention_kwargs: Optional[Dict[str, Any]] = None,\n", " guidance_rescale: float = 0.0,\n", " original_size: Optional[Tuple[int, int]] = None,\n", " crops_coords_top_left: Tuple[int, int] = (0, 0),\n", " target_size: Optional[Tuple[int, int]] = None,\n", " negative_original_size: Optional[Tuple[int, int]] = None,\n", " negative_crops_coords_top_left: Tuple[int, int] = (0, 0),\n", " negative_target_size: Optional[Tuple[int, int]] = None,\n", " network=None,\n", " start_noise=None,\n", " scale=None,\n", " unet=None,\n", "):\n", " # Default image dimensions to the model's default sample size\n", " height = height or self.default_sample_size * self.vae_scale_factor\n", " width = width or self.default_sample_size * self.vae_scale_factor\n", "\n", " original_size = original_size or (height, width)\n", " target_size = target_size or (height, width)\n", "\n", " # Validate inputs and raise an error if they are incorrect\n", " self.check_inputs(\n", " prompt,\n", " prompt_2,\n", " height,\n", " width,\n", " callback_steps,\n", " negative_prompt,\n", " negative_prompt_2,\n", " prompt_embeds,\n", " negative_prompt_embeds,\n", " pooled_prompt_embeds,\n", " negative_pooled_prompt_embeds,\n", " )\n", "\n", " # Determine batch size based on the input prompts\n", " if prompt is not None and isinstance(prompt, str):\n", " batch_size = 1\n", " elif prompt is not None and isinstance(prompt, list):\n", " batch_size = len(prompt)\n", " else:\n", " batch_size = prompt_embeds.shape[0]\n", "\n", " device = self._execution_device\n", "\n", " # Determine whether to apply classifier-free guidance\n", " do_classifier_free_guidance = guidance_scale > 1.0\n", "\n", " # Encode the input prompt\n", " text_encoder_lora_scale = (\n", " cross_attention_kwargs.get(\"scale\", None) if cross_attention_kwargs is not None else None\n", " )\n", " (\n", " prompt_embeds,\n", " negative_prompt_embeds,\n", " pooled_prompt_embeds,\n", " negative_pooled_prompt_embeds,\n", " ) = self.encode_prompt(\n", " prompt=prompt,\n", " prompt_2=prompt_2,\n", " device=device,\n", " num_images_per_prompt=num_images_per_prompt,\n", " do_classifier_free_guidance=do_classifier_free_guidance,\n", " negative_prompt=negative_prompt,\n", " negative_prompt_2=negative_prompt_2,\n", " prompt_embeds=prompt_embeds,\n", " negative_prompt_embeds=negative_prompt_embeds,\n", " pooled_prompt_embeds=pooled_prompt_embeds,\n", " negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,\n", " lora_scale=text_encoder_lora_scale,\n", " )\n", "\n", " # Prepare timesteps for the denoising process\n", " self.scheduler.set_timesteps(num_inference_steps, device=device)\n", " timesteps = self.scheduler.timesteps\n", "\n", " # Prepare latent variables for image generation\n", " num_channels_latents = unet.config.in_channels\n", " latents = self.prepare_latents(\n", " batch_size * num_images_per_prompt,\n", " num_channels_latents,\n", " height,\n", " width,\n", " prompt_embeds.dtype,\n", " device,\n", " generator,\n", " latents,\n", " )\n", "\n", " # Prepare additional arguments for the diffusion process\n", " extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)\n", "\n", " # Prepare additional time embeddings and IDs\n", " add_text_embeds = pooled_prompt_embeds\n", " add_time_ids = self._get_add_time_ids(\n", " original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype\n", " )\n", " if negative_original_size is not None and negative_target_size is not None:\n", " negative_add_time_ids = self._get_add_time_ids(\n", " negative_original_size,\n", " negative_crops_coords_top_left,\n", " negative_target_size,\n", " dtype=prompt_embeds.dtype,\n", " )\n", " else:\n", " negative_add_time_ids = add_time_ids\n", "\n", " if do_classifier_free_guidance:\n", " prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)\n", " add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)\n", " add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)\n", " \n", " # Move embeddings to the specified device (e.g., GPU)\n", " prompt_embeds = prompt_embeds.to(device)\n", " add_text_embeds = add_text_embeds.to(device)\n", " add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)\n", "\n", " num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)\n", "\n", " # Apply denoising_end if specified, adjusting the number of inference steps accordingly\n", " if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1:\n", " discrete_timestep_cutoff = int(\n", " round(\n", " self.scheduler.config.num_train_timesteps\n", " - (denoising_end * self.scheduler.config.num_train_timesteps)\n", " )\n", " )\n", " num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))\n", " timesteps = timesteps[:num_inference_steps]\n", " latents = latents.to(unet.dtype)\n", "\n", " # Begin the denoising process with a progress bar to track the steps\n", " with self.progress_bar(total=num_inference_steps) as progress_bar:\n", " for i, t in enumerate(timesteps):\n", " if t > start_noise:\n", " network.set_lora_slider(scale=0)\n", " else:\n", " network.set_lora_slider(scale=scale)\n", "\n", " # Prepare the latent model input, applying guidance if necessary\n", " latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n", " latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n", "\n", " # Predict the noise residual using the UNet model\n", " added_cond_kwargs = {\"text_embeds\": add_text_embeds, \"time_ids\": add_time_ids}\n", " with network:\n", " noise_pred = unet(\n", " latent_model_input,\n", " t,\n", " encoder_hidden_states=prompt_embeds,\n", " cross_attention_kwargs=cross_attention_kwargs,\n", " added_cond_kwargs=added_cond_kwargs,\n", " return_dict=False,\n", " )[0]\n", "\n", " # Apply classifier-free guidance if enabled\n", " if do_classifier_free_guidance:\n", " noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n", " noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n", "\n", " # Apply guidance rescaling if specified\n", " if do_classifier_free_guidance and guidance_rescale > 0.0:\n", " noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)\n", "\n", " # Compute the previous noisy sample (x_t -> x_t-1)\n", " latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]\n", "\n", " # Update the progress bar and call the callback function, if provided\n", " if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n", " progress_bar.update()\n", " if callback is not None and i % callback_steps == 0:\n", " callback(i, t, latents)\n", "\n", " # Decode the latent variable back to an image\n", " if not output_type == \"latent\":\n", " # Ensure VAE is in float32 mode to avoid overflow issues\n", " needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast\n", "\n", " if needs_upcasting:\n", " self.upcast_vae() # Temporarily cast VAE to float32\n", " latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)\n", "\n", " image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]\n", "\n", " # Cast back to float16 if necessary\n", " if needs_upcasting:\n", " self.vae.to(dtype=torch.float16)\n", " else:\n", " image = latents\n", "\n", " # Apply watermark if available and specified\n", " if not output_type == \"latent\":\n", " if self.watermark is not None:\n", " image = self.watermark.apply_watermark(image)\n", "\n", " # Process the image to the specified output type (e.g., PIL image)\n", " image = self.image_processor.postprocess(image, output_type=output_type)\n", "\n", " # Optionally return the output as a dictionary\n", " if not return_dict:\n", " return (image,)\n", "\n", " return StableDiffusionXLPipelineOutput(images=image)\n", "\n", "# Set the device to CUDA (GPU)\n", "device = 'cuda'\n", "StableDiffusionXLPipeline.__call__ = call # Override the call method of the StableDiffusionXLPipeline\n", "\n", "# Load the CLIP model and processor for evaluating image-text alignment\n", "clip_model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\").to(device)\n", "clip_processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n", "\n", "# Load LPIPS calculator for perceptual similarity evaluation\n", "lpips_alex = lpips.LPIPS(net='alex').to(device)\n", "\n", "def compute_clip_score(prompt, image):\n", " \"\"\"Calculate the CLIP score for a given prompt and image.\"\"\"\n", " inputs = clip_processor(text=prompt, images=image, return_tensors=\"pt\", padding=True).to(device)\n", " outputs = clip_model(**inputs)\n", " return outputs.logits_per_image.item()\n", "\n", "def compute_lpips_score(image1, image2):\n", " \"\"\"Calculate the LPIPS score between two images.\"\"\"\n", " image1 = image1.unsqueeze(0).to(device)\n", " image2 = image2.unsqueeze(0).to(device)\n", " return lpips_alex(image1, image2).item()\n", "\n", "lora_weights = [\n", " 'muscular.pt' # Path to the LoRA weight file\n", "]\n", "\n", "prompts = [\n", " 'A realistic photograph of a person, bokeh, blurred background, 8k',\n", "]\n", "\n", "weight_dtype = torch.float16\n", "start_noise = 700\n", "num_images_per_prompt = 1\n", "scales = [0, 1, 2, 3] # Different scales for the LoRA slider\n", "\n", "all_clip_scores = []\n", "all_lpips_scores = []\n", "\n", "for prompt in prompts:\n", " seed = random.randint(0, 2**15) # Random seed for reproducibility\n", " for lora_weight in lora_weights:\n", " # Load the Stable Diffusion XL model pipeline\n", " pipe = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=weight_dtype)\n", " pipe = pipe.to(device)\n", " unet = pipe.unet\n", "\n", " # Determine the training method based on the LoRA weight file name\n", " if 'full' in lora_weight:\n", " train_method = 'full'\n", " elif 'noxattn' in lora_weight:\n", " train_method = 'noxattn'\n", " else:\n", " train_method = 'noxattn'\n", "\n", " network_type = \"c3lier\"\n", " if train_method == 'xattn':\n", " network_type = 'lierla'\n", "\n", " modules = DEFAULT_TARGET_REPLACE\n", " if network_type == \"c3lier\":\n", " modules += UNET_TARGET_REPLACE_MODULE_CONV\n", "\n", " model_name = lora_weight\n", " name = os.path.basename(model_name)\n", " rank = 4\n", " alpha = 1\n", " if 'rank4' in lora_weight:\n", " rank = 4\n", " if 'rank8' in lora_weight:\n", " rank = 8\n", " if 'alpha1' in lora_weight:\n", " alpha = 1.0\n", "\n", " # Load the LoRA network with the specified configuration\n", " network = LoRANetwork(\n", " unet,\n", " rank=rank,\n", " multiplier=1.0,\n", " alpha=alpha,\n", " train_method=train_method,\n", " ).to(device, dtype=weight_dtype)\n", " network.load_state_dict(torch.load(lora_weight))\n", "\n", " image_list = []\n", " clip_scores = []\n", " lpips_scores = []\n", " for scale in scales:\n", " generator = torch.manual_seed(seed) # Set the random seed for the generator\n", " images = pipe(prompt, num_images_per_prompt=1, num_inference_steps=50, generator=generator, network=network, start_noise=start_noise, scale=scale, unet=unet).images[0]\n", " image_list.append(images)\n", "\n", " # Calculate CLIP score for the generated image\n", " clip_score = compute_clip_score(prompt, images)\n", " clip_scores.append(clip_score)\n", "\n", " # Calculate LPIPS score between consecutive images\n", " if len(image_list) > 1:\n", " lpips_score = compute_lpips_score(transforms.ToTensor()(image_list[-1]), transforms.ToTensor()(image_list[-2]))\n", " lpips_scores.append(lpips_score)\n", " else:\n", " lpips_scores.append(0) # The first image has no previous image to compare\n", "\n", " all_clip_scores.append((prompt, lora_weight, clip_scores))\n", " all_lpips_scores.append((prompt, lora_weight, lpips_scores))\n", "\n", " # Clear up memory and release resources\n", " del unet, network, pipe\n", " unet = None\n", " network = None\n", " pipe = None\n", " torch.cuda.empty_cache()\n", " flush()\n", "\n", " # Display the generated images with titles indicating the scale value\n", " fig, ax = plt.subplots(1, len(image_list), figsize=(20, 4))\n", " for i, a in enumerate(ax):\n", " a.imshow(image_list[i])\n", " a.set_title(f\"{scales[i]}\", fontsize=15)\n", " a.axis('off')\n", "\n", " plt.suptitle(f'{os.path.basename(lora_weight)}', fontsize=20)\n", " plt.show()\n", "\n", "# Store the results in DataFrames and display them\n", "clip_results = []\n", "lpips_results = []\n", "\n", "# Organize CLIP and LPIPS scores into a structured format\n", "for i, (prompt, lora_weight, clip_scores) in enumerate(all_clip_scores):\n", " lpips_scores = all_lpips_scores[i][2]\n", " for scale, (clip_score, lpips_score) in enumerate(zip(clip_scores, lpips_scores)):\n", " clip_results.append((prompt, lora_weight, scale, clip_score))\n", " lpips_results.append((prompt, lora_weight, scale, lpips_score))\n", "\n", "# Create DataFrames from the results for easy analysis\n", "clip_df = pd.DataFrame(clip_results, columns=[\"Prompt\", \"LoRA Weight\", \"Scale\", \"CLIP Score\"])\n", "lpips_df = pd.DataFrame(lpips_results, columns=[\"Prompt\", \"LoRA Weight\", \"Scale\", \"LPIPS Score\"])\n", "\n", "# Print the DataFrame with CLIP scores\n", "print(\"CLIP Scores:\")\n", "print(clip_df)\n", "\n", "# Print the DataFrame with LPIPS scores\n", "print(\"LPIPS Scores:\")\n", "print(lpips_df)\n", "\n", "# Save the CLIP and LPIPS results to CSV files\n", "clip_df.to_csv(\"clip_scores.csv\", index=False)\n", "lpips_df.to_csv(\"lpips_scores.csv\", index=False)\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.5" } }, "nbformat": 4, "nbformat_minor": 5 }