Spaces:
Sleeping
Sleeping
| 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 | |
| blue_loss_scale = 200 | |
| num_inference_steps = 50 | |
| elastic_transformer = T.ElasticTransform(alpha=550.0, sigma=5.0) | |
| pretrained_model_name_or_path = "segmind/tiny-sd" | |
| pipe = DiffusionPipeline.from_pretrained( | |
| pretrained_model_name_or_path, | |
| #low_cpu_mem_usage = True | |
| ).to(torch_device) | |
| pipe.load_textual_inversion("sd-concepts-library/dreams") | |
| pipe.load_textual_inversion("sd-concepts-library/midjourney-style") | |
| pipe.load_textual_inversion("sd-concepts-library/moebius") | |
| pipe.load_textual_inversion("sd-concepts-library/style-of-marc-allante") | |
| pipe.load_textual_inversion("sd-concepts-library/wlop-style") | |
| concepts_mapping = { | |
| "Dream": '<meeg>', "Midjourney":'<midjourney-style>', | |
| "Marc Allante": '<Marc_Allante>', "Moebius": '<moebius>', | |
| "Wlop": '<wlop-style>' | |
| } | |
| def image_loss(images, method='elastic'): | |
| # elastic loss | |
| if method == 'elastic': | |
| transformed_imgs = elastic_transformer(images) | |
| error = torch.abs(transformed_imgs - images).mean() | |
| # symmetry loss - Flip the image along the width | |
| elif method == "symmetry": | |
| flipped_image = torch.flip(images, [3]) | |
| error = F.mse_loss(images, flipped_image) | |
| # saturation loss | |
| elif method == 'saturation': | |
| transformed_imgs = T.functional.adjust_saturation(images,saturation_factor = 10) | |
| error = torch.abs(transformed_imgs - images).mean() | |
| # blue loss | |
| elif method == 'blue': | |
| error = torch.abs(images[:,2] - 0.9).mean() # [:,2] -> all images in batch, only the blue channel | |
| return error | |
| HTML_TEMPLATE = """ | |
| <style> | |
| body { | |
| background: linear-gradient(135deg, #f5f7fa, #c3cfe2); | |
| } | |
| #app-header { | |
| text-align: center; | |
| background: rgba(255, 255, 255, 0.8); /* Semi-transparent white */ | |
| padding: 20px; | |
| border-radius: 10px; | |
| box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); | |
| position: relative; /* To position the artifacts */ | |
| } | |
| #app-header h1 { | |
| color: #4CAF50; | |
| font-size: 2em; | |
| margin-bottom: 10px; | |
| } | |
| .concept { | |
| position: relative; | |
| transition: transform 0.3s; | |
| } | |
| .concept:hover { | |
| transform: scale(1.1); | |
| } | |
| .concept img { | |
| width: 100px; | |
| border-radius: 10px; | |
| box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); | |
| } | |
| .concept-description { | |
| position: absolute; | |
| bottom: -30px; | |
| left: 50%; | |
| transform: translateX(-50%); | |
| background-color: #4CAF50; | |
| color: white; | |
| padding: 5px 10px; | |
| border-radius: 5px; | |
| opacity: 0; | |
| transition: opacity 0.3s; | |
| } | |
| .concept:hover .concept-description { | |
| opacity: 1; | |
| } | |
| /* Artifacts */ | |
| .artifact { | |
| position: absolute; | |
| background: rgba(76, 175, 80, 0.1); /* Semi-transparent green */ | |
| border-radius: 50%; /* Make it circular */ | |
| } | |
| .artifact.large { | |
| width: 300px; | |
| height: 300px; | |
| top: -50px; | |
| left: -150px; | |
| } | |
| .artifact.medium { | |
| width: 200px; | |
| height: 200px; | |
| bottom: -50px; | |
| right: -100px; | |
| } | |
| .artifact.small { | |
| width: 100px; | |
| height: 100px; | |
| top: 50%; | |
| left: 50%; | |
| transform: translate(-50%, -50%); | |
| } | |
| </style> | |
| <div id="app-header"> | |
| <!-- Artifacts --> | |
| <div class="artifact large"></div> | |
| <div class="artifact medium"></div> | |
| <div class="artifact small"></div> | |
| <!-- Content --> | |
| <h1>Art Generator</h1> | |
| <p>Generate new art in five different styles by providing a prompt.</p> | |
| <div style="display: flex; justify-content: center; gap: 20px; margin-top: 20px;"> | |
| <div class="concept"> | |
| <img src="https://github.com/Delve-ERAV1/S20/assets/11761529/30ac92f8-fc62-4aab-9221-043865c6fe7c" alt="Midjourney"> | |
| <div class="concept-description">Midjourney Style</div> | |
| </div> | |
| <div class="concept"> | |
| <img src="https://github.com/Delve-ERAV1/S20/assets/11761529/54c9a61e-df9f-4054-835b-ec2c6ba5916c" alt="Dreams"> | |
| <div class="concept-description">Dreams Style</div> | |
| </div> | |
| <div class="concept"> | |
| <img src="https://github.com/Delve-ERAV1/S20/assets/11761529/2f37e402-15d1-4a74-ba85-bb1566da930e" alt="Moebius"> | |
| <div class="concept-description">Moebius Style</div> | |
| </div> | |
| <div class="concept"> | |
| <img src="https://github.com/Delve-ERAV1/S20/assets/11761529/f838e767-ac20-4996-b5be-65c61b365ce0" alt="Allante"> | |
| <div class="concept-description">Hong Kong born artist inspired by western and eastern influences</div> | |
| </div> | |
| <div class="concept"> | |
| <img src="https://github.com/Delve-ERAV1/S20/assets/11761529/9958140a-1b62-4972-83ca-85b023e3863f" alt="Wlop"> | |
| <div class="concept-description">WLOP (Born 1987) is known for Digital Art (NFTs)</div> | |
| </div> | |
| </div> | |
| </div> | |
| """ | |
| def get_examples(): | |
| examples = [ | |
| ['A powerful man in dreadlocks', 'Dream', 'Symmetry', 45], | |
| ['World Peace', 'Marc Allante', 'Saturation', 147], | |
| ['Storm trooper in the desert, dramatic lighting, high-detail', 'Moebius', 'Elastic', 28], | |
| ['Delicious Italian pizza on a table, a window in the background overlooking a city skyline', 'Wlop', 'Blue', 50], | |
| ] | |
| return(examples) | |
| def latents_to_pil(latents): | |
| # bath of latents -> list of images | |
| latents = (1 / 0.18215) * latents | |
| with torch.no_grad(): | |
| image = pipe.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 generate_art(prompt, concept, method, seed): | |
| prompt = f"{prompt} in the style of {concepts_mapping[concept]}" | |
| img_no_loss = latents_to_pil(generate_image(prompt, method, seed)) | |
| img_loss = latents_to_pil(generate_image(prompt, method, seed, loss_apply=True)) | |
| return([img_no_loss, img_loss]) | |
| def generate_image(prompt, method, seed, loss_apply=False): | |
| generator = torch.manual_seed(seed) | |
| batch_size = 1 | |
| method = method.lower() | |
| # scheduler | |
| scheduler = LMSDiscreteScheduler(beta_start = 0.00085, beta_end = 0.012, beta_schedule = "scaled_linear", num_train_timesteps = 1000) | |
| scheduler.set_timesteps(50) | |
| scheduler.timesteps = scheduler.timesteps.to(torch.float32) | |
| # text embeddings of the prompt | |
| text_input = pipe.tokenizer([prompt], padding='max_length', max_length = pipe.tokenizer.model_max_length, truncation= True, return_tensors="pt") | |
| input_ids = text_input.input_ids.to(torch_device) | |
| with torch.no_grad(): | |
| text_embeddings = pipe.text_encoder(text_input.input_ids.to(torch_device))[0] | |
| max_length = text_input.input_ids.shape[-1] | |
| uncond_input = pipe.tokenizer( | |
| [""] * 1, padding="max_length", max_length= max_length, return_tensors="pt" | |
| ) | |
| with torch.no_grad(): | |
| uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(torch_device))[0] | |
| text_embeddings = torch.cat([uncond_embeddings,text_embeddings]) | |
| # random latent | |
| latents = torch.randn( | |
| (batch_size, pipe.unet.config.in_channels, height// 8, width //8), | |
| generator = generator, | |
| ) | |
| 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 = pipe.unet(latent_model_input, 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_apply and i%5 == 0: | |
| latents = latents.detach().requires_grad_() | |
| latents_x0 = latents - sigma * noise_pred | |
| # use vae to decode the image | |
| denoised_images = pipe.vae.decode((1/ 0.18215) * latents_x0).sample / 2 + 0.5 # range(0,1) | |
| loss = image_loss(denoised_images, method) * blue_loss_scale | |
| 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 |