Gosula/Stable_diffusion_model

1from base64 import b64encode2 3import torch4import numpy as np5from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel6from huggingface_hub import notebook_login7import torch.nn.functional as F8# For video display:9from IPython.display import HTML10from matplotlib import pyplot as plt11from pathlib import Path12from PIL import Image13from torch import autocast14from torchvision import transforms as tfms15from tqdm.auto import tqdm16from transformers import CLIPTextModel, CLIPTokenizer, logging17import os18from device import torch_device,vae,text_encoder,unet,tokenizer,scheduler,token_emb_layer,pos_emb_layer,position_embeddings19 20 21 22# Supress some unnecessary warnings when loading the CLIPTextModel23logging.set_verbosity_error()24 25def pil_to_latent(input_im):26    # Single image -> single latent in a batch (so size 1, 4, 64, 64)27    with torch.no_grad():28        latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling29    return 0.18215 * latent.latent_dist.sample()30 31 32def latents_to_pil(latents):33    # batch of latents -> list of images34    latents = (1 / 0.18215) * latents35    with torch.no_grad():36        image = vae.decode(latents).sample37    image = (image / 2 + 0.5).clamp(0, 1)38    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()39    images = (image * 255).round().astype("uint8")40    pil_images = [Image.fromarray(image) for image in images]41    return pil_images42 43 44def set_timesteps(scheduler, num_inference_steps):45    scheduler.set_timesteps(num_inference_steps)46    scheduler.timesteps = scheduler.timesteps.to(torch.float32)47 48 49def orange_loss(image):50    # Convert the image to a NumPy array51    #image = image.float()  # Convert to a more standard data type (float32)52    #image_np = image.detach().cpu().numpy()  # Use .detach() and .cpu() to ensure compatibility53 54    # Extract the orange channel (e.g., Red and Green channels)55    orange_channel = image[:, 0, :, :] + image[:, 1, :, :]56 57    # Calculate the mean intensity of the orange channel58    #orange_mean = np.mean(orange_channel)59 60    # Define the target mean intensity you desire61    target_mean = 0.8  # Replace with your desired mean intensity62 63    # Calculate the loss based on the squared difference from the target64    loss = torch.abs(orange_channel- target_mean).mean()65 66    # Convert the loss to a PyTorch tensor67    #loss = torch.tensor(loss, dtype=image.dtype)68 69    return loss