app.py

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()