Upload 6 files

config.py

@@ -0,0 +1,20 @@
+import torch
+
+DEVICE = "mps"
+# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+HEIGHT = 512
+WIDTH = 512
+GUIDANCE_SCALE = 8
+LOSS_SCALE = 200
+NUM_INFERENCE_STEPS = 50
+BATCH_SIZE = 1
+
+SEEDS = [2000,2000,500,600,100]
+STABLE_DIFFUSION_MODEL = "CompVis/stable-diffusion-v1-4"
+STABLE_DIFUSION_CONCEPTS = ['<meeg>', '<midjourney-style>', '<moebius>', ' <Marc_Allante>', '<wlop-style>']
+
+#LMS DSCRETE SCHEDULER
+BETA_START = 0.00085
+BETA_END = 0.012
+BETA_SCHEDULE = "scaled_linear"
+NUM_TRAIN_TIMESTEPS = 1000

diffusion_loss.py

@@ -0,0 +1,24 @@
+import torch
+import torchvision.transforms as T
+import torch.nn.functional as F
+
+def blue_channel(images):
+    error = torch.abs(images[:,2] - 0.9).mean() 
+    return error
+
+def elastic_transform(images):
+    elastic_transformer = T.ElasticTransform(alpha=550.0,sigma=5.0)
+    transformed_imgs = elastic_transformer(images)
+    error = torch.abs(transformed_imgs - images).mean()
+    return error
+
+def symmetry(images):
+    flipped_image = torch.flip(images, [3])
+    error = F.mse_loss(images, flipped_image)
+    print("Loss Calculated for the Symmetry : ", error)
+    return error
+
+def saturation(images):
+    transformed_imgs = T.functional.adjust_saturation(images,saturation_factor = 10)
+    error = torch.abs(transformed_imgs - images).mean()
+    return error

image_generator.py

@@ -0,0 +1,73 @@
+import torch
+from tqdm.auto import tqdm
+from diffusers import LMSDiscreteScheduler
+
+import config
+
+def construct_text_embeddings(pipe, prompt):
+    text_input = pipe.tokenizer(prompt, padding='max_length', 
+                                max_length = pipe.tokenizer.model_max_length, truncation= True, 
+                                return_tensors="pt")
+    uncond_input = pipe.tokenizer([""] * config.BATCH_SIZE, padding="max_length", 
+                                  max_length= text_input.input_ids.shape[-1], 
+                                  return_tensors="pt")
+    with torch.no_grad():
+        text_input_embeddings = pipe.text_encoder(text_input.input_ids.to(config.DEVICE))[0]
+    with torch.no_grad():
+        uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(config.DEVICE))[0]
+    
+    text_embeddings = torch.cat([uncond_embeddings, text_input_embeddings])
+    return text_embeddings
+
+def initialize_latent(seed_number, pipe, scheduler):
+    generator = torch.manual_seed(seed_number)
+    latent = torch.randn((config.BATCH_SIZE, pipe.unet.config.in_channels, 
+                           config.HEIGHT//8, config.WIDTH//8), 
+                           generator = generator).to(torch.float16)
+    latent = latent.to(config.DEVICE)
+    latent = latent * scheduler.init_noise_sigma
+    return latent
+
+def run_prediction(pipe, text_embeddings, scheduler, latent, loss_function=None):
+    for i, t in tqdm(enumerate(scheduler.timesteps), total = len(scheduler.timesteps)):
+        latent_model_input = torch.cat([latent] * 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.to(torch.float16), t, encoder_hidden_states=text_embeddings)["sample"]
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + config.GUIDANCE_SCALE * (noise_pred_text - noise_pred_uncond)
+
+        if loss_function and i%5 == 0:
+            latent = latent.detach().requires_grad_()
+            latent_x0 = latent - sigma * noise_pred
+
+            denoised_images = pipe.vae.decode((1/ 0.18215) * latent_x0).sample / 2 + 0.5 # range(0,1)
+
+            loss = loss_function(denoised_images) * config.LOSS_SCALE
+            print(f"loss {loss}")
+
+            cond_grad = torch.autograd.grad(loss, latent)[0]
+            latent = latent.detach() - cond_grad * sigma**2
+        
+        latent = scheduler.step(noise_pred,t, latent).prev_sample
+
+    return latent
+
+def generate_images(pipe, seed_number, prompt, loss_function=None):
+
+    scheduler = LMSDiscreteScheduler(beta_start = 0.00085, 
+                                     beta_end = 0.012, 
+                                     beta_schedule = "scaled_linear", 
+                                     num_train_timesteps = 1000)
+    scheduler.set_timesteps(config.NUM_INFERENCE_STEPS)
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32)
+
+    text_embeddings = construct_text_embeddings(pipe=pipe, prompt=prompt)
+    latent = initialize_latent(seed_number=seed_number, pipe=pipe, scheduler=scheduler)
+    latent = run_prediction(pipe=pipe, text_embeddings=text_embeddings, 
+                            scheduler=scheduler, latent=latent, 
+                            loss_function=loss_function)
+
+    return latent

model.py

@@ -0,0 +1,16 @@
+import torch
+from diffusers import DiffusionPipeline
+import config
+
+def initialize_diffusion_model():
+    pretrained_model_name_or_path = config.STABLE_DIFFUSION_MODEL
+    pipe = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path,
+                                             torch_dtype=torch.float16).to(config.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")
+
+    return pipe

prediction.py

@@ -0,0 +1,24 @@
+import torch
+import gc
+
+import utils
+import model
+import config
+import image_generator as generator
+
+def predict(prompt, pipe, loss_function=None):
+    latents = []
+    
+    for seed_number, sd_concept in zip(config.SEEDS, config.STABLE_DIFUSION_CONCEPTS):
+        torch.mps.empty_cache()
+        gc.collect()
+        torch.mps.empty_cache()
+        
+        prompt = [f'{prompt} {sd_concept}']
+        latent = generator.generate_images(pipe=pipe, seed_number=seed_number, prompt=prompt, loss_function=loss_function)
+        latents.append(latent)
+    
+    latents = torch.vstack(latents)
+    images = utils.convert_latents_to_pil_images(pipe=pipe, latents=latents)
+    grid = utils.populate_image_grid(images, 1, len(latents))
+    return grid

utils.py

@@ -0,0 +1,24 @@
+import torch
+from PIL import Image
+from diffusers import LMSDiscreteScheduler
+
+import config
+
+def convert_latents_to_pil_images(pipe, latents):
+    latents = (1 / 0.18215) * latents
+    with torch.no_grad():
+        image = pipe.vae.decode(latents).sample
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+    images = (image * 255).round().astype("uint8")
+    pil_images = [Image.fromarray(image) for image in images]
+    return pil_images
+
+def populate_image_grid(imgs, rows, cols):
+    assert len(imgs) == rows*cols
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid