import os
import torch
import gradio as gr
from tqdm import tqdm
from PIL import Image
import torch.nn.functional as F
from torchvision import transforms as tfms
from transformers import CLIPTextModel, CLIPTokenizer, logging
from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel, DiffusionPipeline
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 pipeline
model_path = "CompVis/stable-diffusion-v1-4"
sd_pipeline = DiffusionPipeline.from_pretrained(
model_path,
low_cpu_mem_usage=True,
torch_dtype=torch.float32
).to(torch_device)
# Load textual inversions
sd_pipeline.load_textual_inversion("sd-concepts-library/illustration-style")
sd_pipeline.load_textual_inversion("sd-concepts-library/line-art")
sd_pipeline.load_textual_inversion("sd-concepts-library/hitokomoru-style-nao")
sd_pipeline.load_textual_inversion("sd-concepts-library/style-of-marc-allante")
sd_pipeline.load_textual_inversion("sd-concepts-library/midjourney-style")
sd_pipeline.load_textual_inversion("sd-concepts-library/hanfu-anime-style")
sd_pipeline.load_textual_inversion("sd-concepts-library/birb-style")
# Update style token dictionary
style_token_dict = {
"Illustration Style": '<illustration-style>',
"Line Art":'<line-art>',
"Hitokomoru Style":'<hitokomoru-style-nao>',
"Marc Allante": '<Marc_Allante>',
"Midjourney":'<midjourney-style>',
"Hanfu Anime": '<hanfu-anime-style>',
"Birb Style": '<birb-style>'
}
def apply_guidance(latents, guidance_method, loss_scale):
if guidance_method == 'Grayscale':
rgb = latents_to_pil(latents)[0]
gray = rgb.convert('L')
gray_latents = pil_to_latent(gray.convert('RGB'))
return latents + (gray_latents - latents) * loss_scale
elif guidance_method == 'Bright':
bright_latents = F.relu(latents) # Simple brightness increase
return latents + (bright_latents - latents) * loss_scale
elif guidance_method == 'Contrast':
mean = latents.mean()
contrast_latents = (latents - mean) * 2 + mean
return latents + (contrast_latents - latents) * loss_scale
elif guidance_method == 'Symmetry':
flipped_latents = torch.flip(latents, [3]) # Flip horizontally
return latents + (flipped_latents - latents) * loss_scale
elif guidance_method == 'Saturation':
rgb = latents_to_pil(latents)[0]
saturated = tfms.functional.adjust_saturation(tfms.ToTensor()(rgb), 2)
saturated_latents = pil_to_latent(tfms.ToPILImage()(saturated))
return latents + (saturated_latents - latents) * loss_scale
else:
return latents
def generate_with_guidance(prompt, num_inference_steps, guidance_scale, seed, guidance_method, loss_scale):
generator = torch.Generator(device=torch_device).manual_seed(seed)
# Get the text embeddings
text_input = sd_pipeline.tokenizer(prompt, padding="max_length", max_length=sd_pipeline.tokenizer.model_max_length, truncation=True, return_tensors="pt")
with torch.no_grad():
text_embeddings = sd_pipeline.text_encoder(text_input.input_ids.to(torch_device))[0]
# Set the timesteps
sd_pipeline.scheduler.set_timesteps(num_inference_steps)
# Prepare latents
latents = torch.randn(
(1, sd_pipeline.unet.in_channels, 64, 64),
generator=generator,
device=torch_device
)
latents = latents * sd_pipeline.scheduler.init_noise_sigma
# Denoising loop
for t in tqdm(sd_pipeline.scheduler.timesteps):
# Expand the latents for classifier-free guidance
latent_model_input = torch.cat([latents] * 2)
latent_model_input = sd_pipeline.scheduler.scale_model_input(latent_model_input, timestep=t)
# Predict the noise residual
with torch.no_grad():
noise_pred = sd_pipeline.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)
# Apply custom guidance
latents = apply_guidance(latents, guidance_method, loss_scale / 10000) # Normalize loss_scale
# Compute the previous noisy sample x_t -> x_t-1
latents = sd_pipeline.scheduler.step(noise_pred, t, latents).prev_sample
# Scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
image = sd_pipeline.vae.decode(latents).sample
# Convert to PIL Image
image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
image = (image * 255).round().astype("uint8")[0]
image = Image.fromarray(image)
return image
def inference(text, style, inference_step, guidance_scale, seed, guidance_method, loss_scale):
prompt = text + " " + style_token_dict[style]
# Generate image with pipeline
image_pipeline = sd_pipeline(
prompt,
num_inference_steps=inference_step,
guidance_scale=guidance_scale,
generator=torch.Generator(device=torch_device).manual_seed(seed)
).images[0]
# Generate image with guidance
image_guide = generate_with_guidance(prompt, inference_step, guidance_scale, seed, guidance_method, loss_scale)
return image_pipeline, image_guide
title = "Generative with Textual Inversion and Guidance"
description = "A Gradio interface to infer Stable Diffusion and generate images with different art styles and guidance methods"
examples = [
["A majestic castle on a floating island", 'Illustration Style', 20, 7.5, 42, 'Grayscale', 200],
["A cyberpunk cityscape at night", 'Midjourney', 25, 8.0, 123, 'Contrast', 300]
]
demo = gr.Interface(inference,
inputs = [gr.Textbox(label="Prompt", type="text"),
gr.Dropdown(label="Style", choices=list(style_token_dict.keys()), value="Illustration Style"),
gr.Slider(1, 50, 10, step = 1, label="Inference steps"),
gr.Slider(1, 10, 7.5, step = 0.1, label="Guidance scale"),
gr.Slider(0, 10000, 42, step = 1, label="Seed"),
gr.Dropdown(label="Guidance method", choices=['Grayscale', 'Bright', 'Contrast',
'Symmetry', 'Saturation'], value="Grayscale"),
gr.Slider(100, 10000, 200, step = 100, label="Loss scale")],
outputs= [gr.Image(width=512, height=512, label="Generated art"),
gr.Image(width=512, height=512, label="Generated art with guidance")],
title=title,
description=description,
examples=examples)
demo.launch()