Spaces:
Running
Running
| import PIL | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from tqdm import tqdm | |
| import torch.nn.functional as F | |
| import torchvision.transforms as T | |
| from diffusers import LMSDiscreteScheduler, DiffusionPipeline | |
| # configurations | |
| torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" | |
| height, width = 512, 512 | |
| guidance_scale = 8 | |
| loss_scale = 200 | |
| num_inference_steps = 50 | |
| model_path = "CompVis/stable-diffusion-v1-4" | |
| sd_pipeline = DiffusionPipeline.from_pretrained( | |
| model_path, | |
| low_cpu_mem_usage = True, | |
| torch_dtype=torch.float32 | |
| ).to(torch_device) | |
| sd_pipeline.load_textual_inversion("sd-concepts-library/illustration-style") | |
| sd_pipeline.load_textual_inversion("sd-concepts-library/line-art") | |
| sd_pipeline.load_textual_inversion("sd-concepts-library/hitokomoru-style-nao") | |
| sd_pipeline.load_textual_inversion("sd-concepts-library/style-of-marc-allante") | |
| sd_pipeline.load_textual_inversion("sd-concepts-library/midjourney-style") | |
| sd_pipeline.load_textual_inversion("sd-concepts-library/hanfu-anime-style") | |
| sd_pipeline.load_textual_inversion("sd-concepts-library/birb-style") | |
| styles_mapping = { | |
| "Illustration Style": '<illustration-style>', "Line Art":'<line-art>', | |
| "Hitokomoru Style":'<hitokomoru-style-nao>', "Marc Allante": '<Marc_Allante>', | |
| "Midjourney":'<midjourney-style>', "Hanfu Anime": '<hanfu-anime-style>', | |
| "Birb Style": '<birb-style>' | |
| } | |
| # Define seeds for all the styles | |
| seed_list = [11, 56, 110, 65, 5, 29, 47] | |
| # Loss Function based on Edge Detection | |
| def edge_detection(image): | |
| channels = image.shape[1] | |
| # Define the kernels for Edge Detection | |
| ed_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0) | |
| ed_y = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0) | |
| # Replicate the Edge detection kernels for each channel | |
| ed_x = ed_x.repeat(channels, 1, 1, 1).to(image.device) | |
| ed_y = ed_y.repeat(channels, 1, 1, 1).to(image.device) | |
| # ed_x = ed_x.to(torch.float16) | |
| # ed_y = ed_y.to(torch.float16) | |
| # Convolve the image with the Edge detection kernels | |
| conv_ed_x = F.conv2d(image, ed_x, padding=1, groups=channels) | |
| conv_ed_y = F.conv2d(image, ed_y, padding=1, groups=channels) | |
| # Combine the x and y gradients after convolution | |
| ed_value = torch.sqrt(conv_ed_x**2 + conv_ed_y**2) | |
| return ed_value | |
| def edge_loss(image): | |
| ed_value = edge_detection(image) | |
| ed_capped = (ed_value > 0.5).to(torch.float32) | |
| return F.mse_loss(ed_value, ed_capped) | |
| def compute_loss(original_image, loss_type): | |
| if loss_type == 'blue': | |
| # blue loss | |
| # [:,2] -> all images in batch, only the blue channel | |
| error = torch.abs(original_image[:,2] - 0.9).mean() | |
| elif loss_type == 'edge': | |
| # edge loss | |
| error = edge_loss(original_image) | |
| elif loss_type == 'contrast': | |
| # RGB to Gray loss | |
| transformed_image = T.functional.adjust_contrast(original_image, contrast_factor = 2) | |
| error = torch.abs(transformed_image - original_image).mean() | |
| elif loss_type == 'brightness': | |
| # brightnesss loss | |
| transformed_image = T.functional.adjust_brightness(original_image, brightness_factor = 2) | |
| error = torch.abs(transformed_image - original_image).mean() | |
| elif loss_type == 'sharpness': | |
| # sharpness loss | |
| transformed_image = T.functional.adjust_sharpness(original_image, sharpness_factor = 2) | |
| error = torch.abs(transformed_image - original_image).mean() | |
| elif loss_type == 'saturation': | |
| # saturation loss | |
| transformed_image = T.functional.adjust_saturation(original_image, saturation_factor = 10) | |
| error = torch.abs(transformed_image - original_image).mean() | |
| else: | |
| print("error. Loss not defined") | |
| return error | |
| def get_examples(): | |
| examples = [ | |
| ['A bird sitting on a tree', 'Midjourney', 'edge', 5], | |
| ['Cats fighting on the road', 'Marc Allante', 'brightness', 65], | |
| ['A mouse with the head of a puppy', 'Hitokomoru Style', 'contrast', 110], | |
| ['A woman with a smiling face in front of an Italian Pizza', 'Hanfu Anime', 'brightness', 29], | |
| ['A campfire (oil on canvas)', 'Birb Style', 'blue', 47], | |
| ] | |
| return(examples) | |
| def latents_to_pil(latents): | |
| # bath of latents -> list of images | |
| latents = (1 / 0.18215) * latents | |
| with torch.no_grad(): | |
| image = sd_pipeline.vae.decode(latents).sample | |
| image = (image / 2 + 0.5).clamp(0, 1) # 0 to 1 | |
| image = image.detach().cpu().permute(0, 2, 3, 1).numpy() | |
| image = (image * 255).round().astype("uint8") | |
| return Image.fromarray(image[0]) | |
| def show_image(prompt, concept, guidance_type): | |
| for idx, sd in enumerate(styles_mapping.keys()): | |
| if(sd == concept): | |
| break | |
| seed = seed_list[idx] | |
| prompt = f"{prompt} in the style of {styles_mapping[sd]}" | |
| styled_image_without_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=False)) | |
| styled_image_with_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=True)) | |
| return([styled_image_without_loss, styled_image_with_loss]) | |
| def generate_image(seed, prompt, loss_type, loss_flag=False): | |
| generator = torch.manual_seed(seed) | |
| batch_size = 1 | |
| # scheduler | |
| scheduler = LMSDiscreteScheduler(beta_start = 0.00085, beta_end = 0.012, beta_schedule = "scaled_linear", num_train_timesteps = 1000) | |
| scheduler.set_timesteps(num_inference_steps) | |
| scheduler.timesteps = scheduler.timesteps.to(torch.float32) | |
| # text embeddings of the prompt | |
| text_input = sd_pipeline.tokenizer(prompt, padding='max_length', max_length = sd_pipeline.tokenizer.model_max_length, truncation= True, return_tensors="pt") | |
| input_ids = text_input.input_ids.to(torch_device) | |
| with torch.no_grad(): | |
| text_embeddings = sd_pipeline.text_encoder(text_input.input_ids.to(torch_device))[0] | |
| max_length = text_input.input_ids.shape[-1] | |
| uncond_input = sd_pipeline.tokenizer( | |
| [""] * batch_size, padding="max_length", max_length= max_length, return_tensors="pt" | |
| ) | |
| with torch.no_grad(): | |
| uncond_embeddings = sd_pipeline.text_encoder(uncond_input.input_ids.to(torch_device))[0] | |
| text_embeddings = torch.cat([uncond_embeddings,text_embeddings]) # shape: 2,77,768 | |
| # random latent | |
| latents = torch.randn( | |
| (batch_size, sd_pipeline.unet.config.in_channels, height// 8, width //8), | |
| generator = generator, | |
| ) .to(torch.float32) | |
| latents = latents.to(torch_device) | |
| latents = latents * scheduler.init_noise_sigma | |
| for i, t in tqdm(enumerate(scheduler.timesteps), total = len(scheduler.timesteps)): | |
| latent_model_input = torch.cat([latents] * 2) | |
| sigma = scheduler.sigmas[i] | |
| latent_model_input = scheduler.scale_model_input(latent_model_input, t) | |
| with torch.no_grad(): | |
| noise_pred = sd_pipeline.unet(latent_model_input.to(torch.float32), t, encoder_hidden_states=text_embeddings)["sample"] | |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) | |
| noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) | |
| if loss_flag and i%5 == 0: | |
| latents = latents.detach().requires_grad_() | |
| # the following line alone does not work, it requires change to reduce step only once | |
| # hence commenting it out | |
| #latents_x0 = scheduler.step(noise_pred,t, latents).pred_original_sample | |
| latents_x0 = latents - sigma * noise_pred | |
| # use vae to decode the image | |
| denoised_images = sd_pipeline.vae.decode((1/ 0.18215) * latents_x0).sample / 2 + 0.5 # range(0,1) | |
| loss = compute_loss(denoised_images, loss_type) * loss_scale | |
| #loss = loss.to(torch.float16) | |
| print(f"{i} loss {loss}") | |
| cond_grad = torch.autograd.grad(loss, latents)[0] | |
| latents = latents.detach() - cond_grad * sigma**2 | |
| latents = scheduler.step(noise_pred,t, latents).prev_sample | |
| return latents |