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": '', "Line Art":'', "Hitokomoru Style":'', "Marc Allante": '', "Midjourney":'', "Hanfu Anime": '', "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()