Create utils.py

utils.py

@@ -0,0 +1,69 @@
+from base64 import b64encode
+
+import torch
+import numpy as np
+from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
+from huggingface_hub import notebook_login
+import torch.nn.functional as F
+# For video display:
+from IPython.display import HTML
+from matplotlib import pyplot as plt
+from pathlib import Path
+from PIL import Image
+from torch import autocast
+from torchvision import transforms as tfms
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+import os
+from device import torch_device,vae,text_encoder,unet,tokenizer,scheduler,token_emb_layer,pos_emb_layer,position_embeddings
+
+
+
+# Supress some unnecessary warnings when loading the CLIPTextModel
+logging.set_verbosity_error()
+
+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):
+    # batch of latents -> list of images
+    latents = (1 / 0.18215) * latents
+    with torch.no_grad():
+        image = 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 set_timesteps(scheduler, num_inference_steps):
+    scheduler.set_timesteps(num_inference_steps)
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32)
+
+
+def orange_loss(image):
+    # Convert the image to a NumPy array
+    #image = image.float()  # Convert to a more standard data type (float32)
+    #image_np = image.detach().cpu().numpy()  # Use .detach() and .cpu() to ensure compatibility
+
+    # Extract the orange channel (e.g., Red and Green channels)
+    orange_channel = image[:, 0, :, :] + image[:, 1, :, :]
+
+    # Calculate the mean intensity of the orange channel
+    #orange_mean = np.mean(orange_channel)
+
+    # Define the target mean intensity you desire
+    target_mean = 0.8  # Replace with your desired mean intensity
+
+    # Calculate the loss based on the squared difference from the target
+    loss = torch.abs(orange_channel- target_mean).mean()
+
+    # Convert the loss to a PyTorch tensor
+    #loss = torch.tensor(loss, dtype=image.dtype)
+
+    return loss