import open_clip import gradio as gr import numpy as np import torch import torchvision from tqdm.auto import tqdm from PIL import Image, ImageColor from torchvision import transforms from diffusers import DDIMScheduler, DDPMPipeline device = ( "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu" ) # Load the pretrained pipeline pipeline_name = "alkzar90/sd-class-ukiyo-e-256" image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device) # Sample some images with a DDIM Scheduler over 40 steps scheduler = DDIMScheduler.from_pretrained(pipeline_name) scheduler.set_timesteps(num_inference_steps=40) # Color guidance #------------------------------------------------------------------------------- # Color guidance function def color_loss(images, target_color=(0.1, 0.9, 0.5)): """Given a target color (R, G, B) return a loss for how far away on average the images' pixels are from that color. Defaults to a light teal: (0.1, 0.9, 0.5)""" target = ( torch.tensor(target_color).to(images.device) * 2 - 1 ) # Map target color to (-1, 1) target = target[ None, :, None, None ] # Get shape right to work with the images (b, c, h, w) error = torch.abs( images - target ).mean() # Mean absolute difference between the image pixels and the target color return error # CLIP guidance #------------------------------------------------------------------------------- clip_model, _, preprocess = open_clip.create_model_and_transforms( "ViT-B-32", pretrained="openai" ) clip_model.to(device) # Transforms to resize and augment an image + normalize to match CLIP's training data tfms = transforms.Compose( [ transforms.RandomResizedCrop(224), # Random CROP each time transforms.RandomAffine( 5 ), # One possible random augmentation: skews the image transforms.RandomHorizontalFlip(), # You can add additional augmentations if you like transforms.Normalize( mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711), ), ] ) # CLIP guidance function def clip_loss(image, text_features): image_features = clip_model.encode_image( tfms(image) ) # Note: applies the above transforms input_normed = torch.nn.functional.normalize(image_features.unsqueeze(1), dim=2) embed_normed = torch.nn.functional.normalize(text_features.unsqueeze(0), dim=2) dists = ( input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2) ) # Squared Great Circle Distance return dists.mean() # Sample generator loop #------------------------------------------------------------------------------- def generate(color, color_loss_scale, num_examples=4, seed=None, prompt=None, prompt_loss_scale=None, prompt_n_cuts=None, inference_steps=50, ): scheduler.set_timesteps(num_inference_steps=inference_steps) if seed: torch.manual_seed(seed) if prompt: text = open_clip.tokenize([prompt]).to(device) with torch.no_grad(), torch.cuda.amp.autocast(): text_features = clip_model.encode_text(text) target_color = ImageColor.getcolor(color, "RGB") # Target color as RGB target_color = [a / 255 for a in target_color] # Rescale from (0, 255) to (0, 1) x = torch.randn(num_examples, 3, 256, 256).to(device) for i, t in tqdm(enumerate(scheduler.timesteps)): model_input = scheduler.scale_model_input(x, t) with torch.no_grad(): noise_pred = image_pipe.unet(model_input, t)["sample"] x = x.detach().requires_grad_() x0 = scheduler.step(noise_pred, t, x).pred_original_sample # color loss loss = color_loss(x0, target_color) * color_loss_scale cond_color_grad = -torch.autograd.grad(loss, x)[0] # Modify x based solely on the color gradient -> x_cond x_cond = x.detach() + cond_color_grad # prompt loss (modify x_cond with cond_prompt_grad) based on # the original x (not modifified previously with cond_color_grad) if prompt: cond_prompt_grad = 0 for cut in range(prompt_n_cuts): # Set requires grad on x x = x.detach().requires_grad_() # Get the predicted x0: x0 = scheduler.step(noise_pred, t, x).pred_original_sample # Calculate loss prompt_loss = clip_loss(x0, text_features) * prompt_loss_scale # Get gradient (scale by n_cuts since we want the average) cond_prompt_grad -= torch.autograd.grad(prompt_loss, x, retain_graph=True)[0] / prompt_n_cuts # Modify x based on this gradient alpha_bar = scheduler.alphas_cumprod[i] x_cond = ( x_cond + cond_prompt_grad * alpha_bar.sqrt() ) # Note the additional scaling factor here! x = scheduler.step(noise_pred, t, x_cond).prev_sample grid = torchvision.utils.make_grid(x, nrow=4) im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5 im = Image.fromarray(np.array(im * 255).astype(np.uint8)) #im.save("test.jpeg") return im # GRADIO Interface #------------------------------------------------------------------------------- TITLE="Ukiyo-e postal generator service 🎴!" DESCRIPTION="This model is a diffusion model for unconditional image generation of Ukiyo-e images ✍ 🎨. \nThe model was train using fine-tuning with the google/ddpm-celebahq-256 pretrain-model and the dataset: https://huggingface.co/datasets/huggan/ukiyoe2photo" CSS = ".output-image, .input-image, .image-preview {height: 250px !important}" # See the gradio docs for the types of inputs and outputs available inputs = [ gr.ColorPicker(label="color (click on the square to pick the color)", value="#DF5C16"), # Add any inputs you need here gr.Slider(label="color_guidance_scale (how strong to blend the color)", minimum=0, maximum=30, value=6.7), gr.Slider(label="num_examples (# images generated)", minimum=2, maximum=12, value=2, step=4), gr.Number(label="seed (reproducibility and experimentation)", value=666), gr.Text(label="Text prompt (optional)", value=None), gr.Slider(label="prompt_guidance_scale (...)", minimum=0, maximum=1000, value=10), gr.Slider(label="prompt_n_cuts", minimum=4, maximum=12, step=4), gr.Slider(label="Number of inference steps (+ steps -> + guidance effect)", minimum=40, maximum=60, value=40, step=1), ] outputs = gr.Image(label="result") # And the minimal interface demo = gr.Interface( fn=generate, inputs=inputs, outputs=outputs, css=CSS, #examples=[ #["#DF5C16", 6.7, 12, 666, None, None, None, 40], #["#C01660", 13.5, 12, 1990, None, None, None, 40], #["#44CCAA", 8.9, 12, 1512, None, None, None, 40], #["#39A291", 5.0, 2, 666, "A sakura tree", 60, 4, 40], #["#0E0907", 0.0, 12, 666, "A big whale in the ocean", 60, 8, 52], #["#19A617", 4.6, 12, 666, "An island with sunset at background", 140, 4, 47], #], title=TITLE, description=DESCRIPTION, ) if __name__ == "__main__": demo.launch(enable_queue=True)