init commit

Marc Allante.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b0496315f14f212535f9350c3dbf05787ac50a78465d4be2f39a1ba373e4968
+size 3819

app.py

@@ -0,0 +1,69 @@
+import gradio as gr
+
+import os
+import torch
+from image_generator import generate_image_per_prompt_style
+
+torch.manual_seed(11)
+
+
+# Set device
+torch_device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+if "mps" == torch_device:
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+# Define Interface
+title = "Generative Art - Stable Diffusion with Styles and additional guidance"
+
+gr_interface = gr.Interface(
+    generate_image_per_prompt_style,
+    inputs=[
+        gr.Textbox("cat running", label="Prompt"),
+        gr.Dropdown(
+            [
+                "illustration_style",
+                "line-art",
+                "hitokomoru-style",
+                "midjourney-style",
+                "hanfu-anime-style",
+                "birb-style",
+                "style-of-marc-allante",
+            ],
+            value="birb-style",
+            label="Pre-trained Styles",
+        ),
+        gr.Dropdown(
+            [
+                "blue_loss",
+                "cosine_loss",
+            ],
+            value="cosine_loss",
+            label="Additional guidance for image generation",
+        ),
+        gr.Textbox("on a city road", label="Additional Prompt"),
+    ],
+    outputs=[
+        gr.Gallery(
+            label="Generated images",
+            show_label=False,
+            elem_id="gallery",
+            columns=[2],
+            rows=[2],
+            object_fit="contain",
+            height="auto",
+        )
+    ],
+    title=title,
+    examples=[
+        ["A flying bird", "illustration_style", "blue_loss", ""],
+        ["cat running", "on a city road", "cosine_loss", ""]
+    ]
+)
+gr_interface.launch(debug=True)
+
+

birb-style.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2e23a8f2d3628ed77acb8151751ecd4efc4017e8da86bc29af10f855ca308d9
+size 3819

hanfu-anime-style.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18ee85c31cff7a0ab35f90af24fbf1a4ab8a9960ab041511e386d5990953e050
+size 3819

hitokomoru-style.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f81a9c575e329e08a24e08f47ae73c5b50dec4bcb557974552549b45e2d1b0d4
+size 3819

illustration_style.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44d65046c071e37f75f31a7a81a34c50a96080e8a3aedc7cda1094dae5d385f0
+size 3819

image_generator.py

@@ -0,0 +1,381 @@
+import os
+from pathlib import Path
+import torch
+from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+
+from utils import load_embedding_bin, set_timesteps, latents_to_pil
+from loss import blue_loss, cosine_loss
+from matplotlib import pyplot as plt
+from pathlib import Path
+
+torch.manual_seed(11)
+logging.set_verbosity_error()
+
+# Set device
+torch_device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+if "mps" == torch_device:
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+# Style embeddings
+STYLE_EMBEDDINGS = {
+    "illustration-style": "illustration_style.bin",
+    "line-art": "line-art.bin",
+    "hitokomoru-style": "hitokomoru-style.bin",
+    "midjourney-style": "midjourney-style.bin",
+    "hanfu-anime-style": "hanfu-anime-style.bin",
+    "birb-style": "birb-style.bin",
+    "style-of-marc-allante": "Marc Allante.bin",
+}
+LOSS = {"blue_loss": blue_loss,
+        "cosine_loss": cosine_loss}
+STYLE_SEEDS = [11, 56, 110, 65, 5, 29, 47]
+# Load the autoencoder model which will be used to decode the latents into image space.
+vae = AutoencoderKL.from_pretrained(
+    "CompVis/stable-diffusion-v1-4", subfolder="vae"
+).to(torch_device)
+#
+# # Load the tokenizer and text encoder to tokenize and encode the text.
+tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(
+    torch_device
+)
+#
+# # The UNet model for generating the latents.
+unet = UNet2DConditionModel.from_pretrained(
+    "CompVis/stable-diffusion-v1-4", subfolder="unet"
+).to(torch_device)
+#
+# # The noise scheduler
+scheduler = LMSDiscreteScheduler(
+    beta_start=0.00085,
+    beta_end=0.012,
+    beta_schedule="scaled_linear",
+    num_train_timesteps=1000,
+)
+
+# vae = vae
+# text_encoder = text_encoder.to(torch_device)
+unet = unet
+token_emb_layer = text_encoder.text_model.embeddings.token_embedding
+pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
+position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
+position_embeddings = pos_emb_layer(position_ids)
+
+
+def build_causal_attention_mask(bsz, seq_len, dtype):
+    # lazily create causal attention mask, with full attention between the vision tokens
+    # pytorch uses additive attention mask; fill with -inf
+    mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype)
+    mask.fill_(torch.tensor(torch.finfo(dtype).min))
+    mask.triu_(1)  # zero out the lower diagonal
+    mask = mask.unsqueeze(1)  # expand mask
+    return mask
+
+
+def get_output_embeds(input_embeddings):
+    # CLIP's text model uses causal mask, so we prepare it here:
+    bsz, seq_len = input_embeddings.shape[:2]
+    causal_attention_mask = build_causal_attention_mask(
+        bsz, seq_len, dtype=input_embeddings.dtype
+    )
+
+    # Getting the output embeddings involves calling the model with passing output_hidden_states=True
+    # so that it doesn't just return the pooled final predictions:
+    encoder_outputs = text_encoder.text_model.encoder(
+        inputs_embeds=input_embeddings,
+        attention_mask=None,  # We aren't using an attention mask so that can be None
+        causal_attention_mask=causal_attention_mask.to(torch_device),
+        output_attentions=None,
+        output_hidden_states=True,  # We want the output embs not the final output
+        return_dict=None,
+    )
+
+    # We're interested in the output hidden state only
+    output = encoder_outputs[0]
+
+    # There is a final layer norm we need to pass these through
+    output = text_encoder.text_model.final_layer_norm(output)
+
+    # And now they're ready!
+    return output
+
+
+# Generating an image with these modified embeddings
+def generate_with_embs(text_embeddings, seed, max_length):
+    height = 512  # default height of Stable Diffusion
+    width = 512  # default width of Stable Diffusion
+    num_inference_steps = 30  # Number of denoising steps
+    guidance_scale = 7.5  # Scale for classifier-free guidance
+    generator = torch.manual_seed(seed)
+    batch_size = 1
+
+    # tokenizer
+    uncond_input = tokenizer(
+        [""] * batch_size,
+        padding="max_length",
+        max_length=max_length,
+        return_tensors="pt",
+    )
+
+    with torch.no_grad():
+        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
+    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+    # Prep Scheduler
+    set_timesteps(scheduler, num_inference_steps)
+
+    # Prep latents
+    # step = " prep_latents "
+    latents = torch.randn(
+        (batch_size, unet.in_channels, height // 8, width // 8),
+        generator=generator,
+    )
+    latents = latents.to(torch_device)
+    latents = latents * scheduler.init_noise_sigma
+
+    # Loop
+    for i, t in tqdm(enumerate(scheduler.timesteps), 
+                     total=len(scheduler.timesteps)):
+        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+        latent_model_input = torch.cat([latents] * 2)
+        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+
+        # predict the noise residual
+        with torch.no_grad():
+            noise_pred = unet(
+                latent_model_input, t, encoder_hidden_states=text_embeddings
+            )["sample"]
+
+        # perform guidance
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+        # compute the previous noisy sample x_t -> x_t-1
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+    return latents_to_pil(latents)[0]
+
+
+def generate_image_from_embeddings(
+        mod_output_embeddings, seed, max_length, 
+        loss_selection, additional_prompt):
+    height = 512
+    width = 512
+    num_inference_steps = 50
+    guidance_scale = 8
+    generator = torch.manual_seed(seed)
+    batch_size = 1
+    if loss_selection == "blue_loss":
+        loss_fn = LOSS["blue_loss"]
+        loss_scale = 120
+    else:
+        loss_fn = LOSS["cosine_loss"](additional_prompt)
+        loss_scale = 20
+
+    # Use the modified_output_embeddings directly
+    text_embeddings = mod_output_embeddings
+
+    uncond_input = tokenizer(
+        [""] * batch_size,
+        padding="max_length",
+        max_length=max_length,
+        return_tensors="pt",
+    )
+    with torch.no_grad():
+        uncond_embeddings = text_encoder(
+            uncond_input.input_ids.to(torch_device))[0]
+    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+    # Prep Scheduler
+    set_timesteps(scheduler, num_inference_steps)
+
+    # Prep latents
+    latents = torch.randn(
+        (batch_size, unet.config.in_channels, height // 8, width // 8),
+        generator=generator,
+    )
+    latents = latents.to(torch_device)
+    latents = latents * scheduler.init_noise_sigma
+
+    # Loop
+    for i, t in tqdm(enumerate(scheduler.timesteps), 
+                     total=len(scheduler.timesteps)):
+        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+        latent_model_input = torch.cat([latents] * 2)
+        sigma = scheduler.sigmas[i]
+        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+
+        # predict the noise residual
+        with torch.no_grad():
+            noise_pred = unet(
+                latent_model_input, t, encoder_hidden_states=text_embeddings
+            )["sample"]
+
+        # perform CFG
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+        #### ADDITIONAL GUIDANCE ###
+        if i % 2 == 0:
+            # Requires grad on the latents
+            latents = latents.detach().requires_grad_()
+
+            # Get the predicted x0:
+            # latents_x0 = latents - sigma * noise_pred
+            latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample
+            scheduler._step_index -= 1
+            # Decode to image space
+            denoised_images = (
+                vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
+            )  # range (0, 1)
+
+            # Calculate loss
+            loss = loss_fn(denoised_images) * loss_scale
+
+            # Occasionally print it out
+            if i % 10 == 0:
+                print(i, "loss:", loss.item())
+
+            # Get gradient
+            cond_grad = torch.autograd.grad(loss, latents)[0]
+
+            # Modify the latents based on this gradient
+            latents = latents.detach() - cond_grad * sigma**2
+
+        # Now step with scheduler
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+    return latents_to_pil(latents)[0]
+
+
+def generate_image_per_style(prompt, style_embed, style_seed, style_embedding_key):
+    modified_output_embeddings = None
+    gen_out_style_image = None
+    max_length = 0
+
+    # Tokenize
+    text_input = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt",
+    )
+    input_ids = text_input.input_ids.to(torch_device)
+
+    # Get token embeddings
+    token_embeddings = token_emb_layer(input_ids)
+
+    replacement_token_embedding = style_embed[style_embedding_key]
+
+    # Insert this into the token embeddings
+    token_embeddings[
+        0, torch.where(input_ids[0] == 6829)[0]
+    ] = replacement_token_embedding.to(torch_device)
+
+    # Combine with pos embs
+    input_embeddings = token_embeddings + position_embeddings
+
+    #  Feed through to get final output embs
+    modified_output_embeddings = get_output_embeds(input_embeddings)
+
+    # And generate an image with this:
+    max_length = text_input.input_ids.shape[-1]
+
+    gen_out_style_image = generate_with_embs(
+        modified_output_embeddings, style_seed, max_length
+    )
+
+    return gen_out_style_image
+
+
+def generate_image_per_loss(
+        prompt, style_embed, style_seed, style_embedding_key,
+        loss, additional_prompt
+    ):
+    gen_out_loss_image = None
+
+    # Tokenize
+    text_input = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt",
+    )
+    input_ids = text_input.input_ids.to(torch_device)
+
+    # Get token embeddings
+    token_embeddings = token_emb_layer(input_ids)
+
+    replacement_token_embedding = style_embed[style_embedding_key].to(torch_device)
+
+    # Insert this into the token embeddings
+    token_embeddings[
+        0, torch.where(input_ids[0] == 6829)[0]
+    ] = replacement_token_embedding
+
+    # Combine with pos embs
+    input_embeddings = token_embeddings + position_embeddings
+    modified_output_embeddings = get_output_embeds(input_embeddings)
+
+    # max_length = tokenizer.model_max_length
+
+    max_length = text_input.input_ids.shape[-1]
+    gen_out_loss_image = generate_image_from_embeddings(
+        modified_output_embeddings, style_seed, max_length,
+        loss, additional_prompt
+    )
+
+    return gen_out_loss_image
+
+
+def generate_image_per_prompt_style(text_in, style_in, 
+                                    loss, additional_prompt):
+    gen_style_image = None
+    gen_loss_image = None
+    STYLE_KEYS = []
+    style_key = ""
+
+    if style_in not in STYLE_EMBEDDINGS:
+        raise ValueError(
+            f"Unknown style: {style_in}. Available styles are: {', '.join(STYLE_EMBEDDINGS.keys())}"
+        )
+
+    STYLE_SEEDS = [32, 64, 128, 16, 8, 96]
+    STYLE_KEYS = list(STYLE_EMBEDDINGS.keys())
+    print(f"prompt: {text_in}")
+    print(f"style: {style_in}")
+
+    idx = STYLE_KEYS.index(style_in)
+    style_file = STYLE_EMBEDDINGS[style_in]
+    print(f"style_file: {style_file}")
+
+    prompt = text_in
+
+    style_seed = STYLE_SEEDS[idx]
+
+    style_key = Path(style_file).stem
+    style_key = style_key.replace("_", "-")
+    print(style_key, STYLE_KEYS, style_file)
+
+    file_path = os.path.join(os.getcwd(), style_file)
+    embedding = load_embedding_bin(file_path)
+    style_key = f"<{style_key}>"
+
+    gen_style_image = generate_image_per_style(prompt, embedding, style_seed, style_key)
+
+    gen_loss_image = generate_image_per_loss(prompt, embedding, style_seed, style_key, loss, additional_prompt)
+
+    return [gen_style_image, gen_loss_image]

line-art.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0528436ec2228c659e0cf1316e713345bc97a3d88294f1a2987a3505d220e770
+size 3819

loss.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms import v2
+from transformers import CLIPTextModel, CLIPTokenizer, \
+ CLIPProcessor, CLIPVisionModelWithProjection, CLIPTextModelWithProjection
+
+import os
+# from image_generator import get_output_embeds, position_embeddings
+
+
+# Set device
+torch_device = "cuda" if torch.cuda.is_available() else "mps" \
+    if torch.backends.mps.is_available() else "cpu"
+
+if "mps" == torch_device: 
+    os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = "1"
+
+# Load the tokenizer and text encoder to tokenize and encode the text.
+clip_model_name = "openai/clip-vit-large-patch14"
+tokenizer = CLIPTokenizer.from_pretrained(clip_model_name)
+text_encoder = CLIPTextModel.from_pretrained(clip_model_name).to(torch_device);
+vision_encoder = CLIPVisionModelWithProjection.from_pretrained(clip_model_name).to(torch_device);
+processor = CLIPProcessor.from_pretrained(clip_model_name)
+
+# # additional textual prompt
+def get_text_embed(prompt = "on a mountain"):
+    inputs = processor(text=prompt,
+                       return_tensors="pt",
+                       padding=True)
+    with torch.no_grad():
+        text_embed = CLIPTextModelWithProjection.from_pretrained(
+            clip_model_name)(**inputs).text_embeds.to(torch_device)
+    return text_embed
+
+# def get_text_embed(prompt = "on a mountain"):
+#     text_input = tokenizer([prompt],
+#                            padding="max_length",
+#                            max_length=tokenizer.model_max_length,
+#                            truncation=True,
+#                            return_tensors="pt")
+#     with torch.no_grad():
+#         text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
+#         input_embeddings = text_embeddings + position_embeddings.to(torch_device)
+#     modified_output_embeddings = get_output_embeds(input_embeddings)
+#     return modified_output_embeddings
+
+class cosine_loss(nn.Module):
+    def __init__(self, prompt) -> None:
+        self.text_embed = get_text_embed(prompt)
+        super().__init__()
+
+    def forward(self, gen_image):
+        gen_image_clamped = gen_image.clamp(0, 1).mul(255)
+        resized_image = v2.Resize(224)(gen_image_clamped)
+        image_embed = vision_encoder(resized_image).image_embeds
+        similarity = F.cosine_similarity(self.text_embed, image_embed, dim=1)
+        loss = 1 - similarity.mean()
+        return loss
+
+def blue_loss(images):
+    # How far are the blue channel values to 0.9:
+    error = torch.abs(images[:,2] - 0.9).mean() # [:,2] -> all images in batch, only the blue channel
+    return error

midjourney-style.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4865a5d2ecd012985940748023fd80e4fd299837f1dccedb85ee83be5bb1f957
+size 3819

utils.py

@@ -0,0 +1,30 @@
+import torch
+from PIL import Image
+from diffusers import AutoencoderKL
+
+vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").to("mps:0")
+
+def pil_to_latent(input_im):
+    # Single image -> single latent in a batch (so size 1, 4, 64, 64)
+    with torch.no_grad():
+        latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling
+    return 0.18215 * latent.latent_dist.sample()
+
+def latents_to_pil(latents, torch_device="mps:0"):
+    # bath of latents -> list of images
+    latents = (1 / 0.18215) * latents
+    with torch.no_grad():
+        image = vae.decode(latents.to(torch_device)).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 load_embedding_bin(path):
+    return torch.load(path)
+
+#  Prep Scheduler
+def set_timesteps(scheduler, num_inference_steps):
+    scheduler.set_timesteps(num_inference_steps)
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32) # minor fix to ensure MPS compatibility, fixed in diffusers PR 3925