Spaces:

ToletiSri
/

TSAI_S20

Runtime error

App Files Files Community

TSAI_S20 / app.py

ToletiSri

Update app.py

3283195 8 months ago

raw history blame contribute delete

No virus

7.95 kB

	import gradio as gr
	from base64 import b64encode

	import numpy
	import torch
	from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel, StableDiffusionPipeline

	# For video display:
	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

	torch.manual_seed(1)

	# Supress some unnecessary warnings when loading the CLIPTextModel
	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"

	model_nm = "CompVis/stable-diffusion-v1-4"

	output_dir="sd-concept-output"
	pipe = StableDiffusionPipeline.from_pretrained(model_nm).to(torch_device)

	# Load the autoencoder model which will be used to decode the latents into image space.
	vae = pipe.vae
	tokenizer = pipe.tokenizer

	# Load the tokenizer and text encoder to tokenize and encode the text.
	#tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16)
	text_encoder =pipe.text_encoder

	# The UNet model for generating the latents.
	unet = pipe.unet

	# The noise scheduler
	scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

	# To the GPU we go!
	vae = vae.to(torch_device)
	text_encoder = text_encoder.to(torch_device)
	unet = unet.to(torch_device);

	pipe.load_textual_inversion("sd-concepts-library/madhubani-art")
	pipe.load_textual_inversion("sd-concepts-library/line-art")
	pipe.load_textual_inversion("sd-concepts-library/cat-toy")
	pipe.load_textual_inversion("sd-concepts-library/concept-art")

	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):
	# bath 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

	# 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

	def saturation_loss(images):
	# Calculate saturation for each pixel in the input image tensor
	max_vals, _ = torch.max(images, dim=1, keepdim=True)
	min_vals, _ = torch.min(images, dim=1, keepdim=True)
	saturation = (max_vals - min_vals) / max_vals.clamp(min=1e-7) # Avoid division by zero

	# Calculate mean saturation across the image
	mean_saturation = torch.mean(saturation, dim=(2, 3)) # Average over width and height

	# Calculate the loss as the negative mean saturation (proportional to saturation)
	#loss = torch.abs(saturation - 0.9).mean()

	return mean_saturation/10000

	def generateImage(prompt, lossScale):
	#prompt = 'a puppy in <cat-toy> style' #@param
	height = 512 # default height of Stable Diffusion
	width = 512 # default width of Stable Diffusion
	num_inference_steps = 200 #@param # Number of denoising steps
	guidance_scale = 8 #@param # Scale for classifier-free guidance
	generator = torch.manual_seed(32) # Seed generator to create the inital latent noise
	batch_size = 1
	saturation_loss_Scale = lossScale #@param

	# Prep text
	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]

	# And the uncond. input as before:
	max_length = text_input.input_ids.shape[-1]
	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.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%5 == 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 = 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 = saturation_loss(denoised_images) * saturation_loss_Scale
	#loss = loss.detach().requires_grad_()
	#print('loss.grad_fn = {}'.format(grad_fn))

	# 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


	custom_loss_image = latents_to_pil(latents)[0]
	return custom_loss_image

	def inference(imgText, style, customLoss="no"):
	prompt = f'a {imgText} in <{style}> style'
	if (customLoss == "yes") :
	outImage = generateImage(prompt, 2)
	return outImage
	else:
	outImage = generateImage(prompt, 0)
	return outImage


	title = "TSAI S20 Assignment: Use a pretrained Sstable Diffusion model and give a demo on its workig"
	description = "A simple Gradio interface that accepts a text and style, and generated an image using stable diffusion pipeline"

	examples = [["puppy","cat-toy","yes"]]

	demo = gr.Interface(
	inference,
	inputs = [gr.Textbox("Enter an image you want to generate"),
	gr.Dropdown(["madhubani-art", "line-art", "cat-toy","concept-art"], label="Choose your style"),
	gr.Radio(["yes", "no"], label="Add custom saturation loss?")
	],
	outputs = [gr.Image(shape=(512, 512), label="Generated Image")],
	title = title,
	description = description,
	examples = examples,
	)
	demo.launch()