First commit

app.py

@@ -0,0 +1,197 @@
+import os
+import torch
+import gradio as gr
+from tqdm import tqdm
+from PIL import Image
+from torchvision import transforms as tfms
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
+
+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 autoencoder model which will be used to decode the latents into image space.
+vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
+
+# Load the tokenizer and text encoder to tokenize and encode the text.
+tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+# The UNet model for generating the latents.
+unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
+
+# The noise scheduler
+scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+
+style_token_dict = {'Concept':'<concept-art>', 'Realistic':'<doose-realistic>', 'Line':'<line-art>',
+                    'Ricky':'<RickyArt>', 'Plane Scape':'<tony-diterlizzi-planescape>'}
+
+# To the GPU we go!
+vae = vae.to(torch_device)
+text_encoder = text_encoder.to(torch_device)
+unet = unet.to(torch_device)
+
+token_emb_layer = text_encoder.text_model.embeddings.token_embedding
+pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
+position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
+position_embeddings = pos_emb_layer(position_ids)
+
+concept_art_embed = torch.load('concept-art.bin')
+doose_s_realistic_art_style_embed = torch.load('doose-s-realistic-art-style.bin')
+line_art_embed = torch.load('line-art.bin')
+rickyart_embed = torch.load('rickyart.bin')
+tony_diterlizzi_s_planescape_art_embed = torch.load('tony-diterlizzi-s-planescape-art.bin')
+
+tokenizer.add_tokens(['<concept-art>', '<doose-realistic>', '<line-art>', '<RickyArt>', '<tony-diterlizzi-planescape>'])
+
+token_emb_layer_with_art = torch.nn.Embedding(49413, 768)
+token_emb_layer_with_art.load_state_dict({'weight': torch.cat((token_emb_layer.state_dict()['weight'],
+                                              concept_art_embed['<concept-art>'].unsqueeze(0).to(torch_device),
+                                              doose_s_realistic_art_style_embed['<doose-realistic>'].unsqueeze(0).to(torch_device),
+                                              line_art_embed['<line-art>'].unsqueeze(0).to(torch_device),
+                                              rickyart_embed['<RickyArt>'].unsqueeze(0).to(torch_device),
+                                              tony_diterlizzi_s_planescape_art_embed['<tony-diterlizzi-planescape>'].unsqueeze(0).to(torch_device)))})
+token_emb_layer_with_art = token_emb_layer_with_art.to(torch_device)
+
+def set_timesteps(scheduler, num_inference_steps):
+    scheduler.set_timesteps(num_inference_steps)
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32)
+
+def pil_to_latent(input_im):
+    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):
+    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
+
+def build_causal_attention_mask(bsz, seq_len, dtype):
+    mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype)
+    mask.fill_(torch.tensor(torch.finfo(dtype).min))  # fill with large negative number (acts like -inf)
+    mask = mask.triu_(1)  # zero out the lower diagonal to enforce causality
+    return mask.unsqueeze(1)  # add a batch dimension
+
+def get_output_embeds(input_embeddings):
+    # CLIP's text model uses causal mask, so we prepare it here:
+    bsz, seq_len = input_embeddings.shape[:2]
+    causal_attention_mask = build_causal_attention_mask(bsz, seq_len, dtype=input_embeddings.dtype)
+
+    # Getting the output embeddings involves calling the model with passing output_hidden_states=True
+    # so that it doesn't just return the pooled final predictions:
+    encoder_outputs = text_encoder.text_model.encoder(
+        inputs_embeds=input_embeddings,
+        attention_mask=None, # We aren't using an attention mask so that can be None
+        causal_attention_mask=causal_attention_mask.to(torch_device),
+        output_attentions=None,
+        output_hidden_states=True, # We want the output embs not the final output
+        return_dict=None,
+    )
+
+    # We're interested in the output hidden state only
+    output = encoder_outputs[0]
+
+    # There is a final layer norm we need to pass these through
+    output = text_encoder.text_model.final_layer_norm(output)
+
+    # And now they're ready!
+    return output
+
+def generate_with_embs(num_inference_steps, guidance_scale, seed, text_input, text_embeddings):
+    height = 512                        # default height of Stable Diffusion
+    width = 512                         # default width of Stable Diffusion
+    generator = torch.manual_seed(seed)   # Seed generator to create the inital latent noise
+    batch_size = 1
+
+    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 guidance
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        # compute the previous noisy sample x_t -> x_t-1
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+    return latents_to_pil(latents)[0]
+
+def inference(text, style, inference_step, guidance_scale, seed):
+    prompt = text + " the style of " + style_token_dict[style]
+
+    # Tokenize
+    text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    input_ids = text_input.input_ids.to(torch_device)
+
+    # Get token embeddings
+    token_embeddings = token_emb_layer_with_art(input_ids)
+
+    # Combine with pos embs
+    input_embeddings = token_embeddings + position_embeddings
+
+    #  Feed through to get final output embs
+    modified_output_embeddings = get_output_embeds(input_embeddings)
+
+    # And generate an image with this:
+    image = generate_with_embs(inference_step, guidance_scale, seed, text_input, modified_output_embeddings)
+
+    return image
+
+title = "Stable Diffusion with Textual Inversion"
+description = "A simple Gradio interface to infer Stable Diffusion and generate images with different art style"
+examples = [["A sweet potato farm", 'Concept', 30, 0.5, 1],
+            ["Sky full of cotton candy", 'Realistic', 30, 1.5, 2],
+            ["Coffin full of jello", 'Line', 30, 2.5, 3],
+            ["Water skiing on a lake", 'Ricky', 30, 3.5, 4],
+            ["Super slippery noodles", 'Plane Scape', 30, 4.5, 5],
+            ["Beautiful sunset", 'Concept', 30, 5.5, 6],
+            ["A glittering gem", 'Realistic', 30, 6.5, 7],
+            ["River rafting", 'Line', 30, 7.5, 8],
+            ["A green tea", 'Ricky', 30, 8.5, 9],
+            ["Three sphered rocks", 'Plane Scape', 30, 9.5, 10]]
+
+demo = gr.Interface(inference, 
+                    inputs = [gr.Textbox(label="Prompt", type="text"),
+                              gr.Dropdown(label="Style", choices=['Concept', 'Realistic', 'Line', 
+                                                                  'Ricky', 'Plane Scape'], value="Concept"), 
+                              gr.Slider(10, 50, 30, step = 10, label="Inference steps"),
+                              gr.Slider(1, 10, 7.5, step = 0.1, label="Guidance scale"),
+                              gr.Slider(0, 10000, 1, step = 1, label="Seed")],
+                    outputs= [gr.Image(width=320, height=320, label="Output SAM")],
+                    title=title,
+                    description=description,
+                    examples=examples)
+
+demo.launch()

concept-art.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd046d1c90c6e58769033de23adadf936e873597b11fed16a07dde7750bd348c
+size 3819

doose-s-realistic-art-style.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a2cfe14214eb4055475b3445420796adfe6aa1bfbc9e1fb7dc62dedd5d71808
+size 3819

line-art.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0528436ec2228c659e0cf1316e713345bc97a3d88294f1a2987a3505d220e770
+size 3819

requirements.txt

@@ -0,0 +1,88 @@
+aiofiles==23.2.1
+altair==5.3.0
+annotated-types==0.7.0
+anyio==4.4.0
+attrs==23.2.0
+certifi==2024.6.2
+charset-normalizer==3.3.2
+click==8.1.7
+colorama==0.4.6
+contourpy==1.2.1
+cycler==0.12.1
+diffusers==0.29.2
+dnspython==2.6.1
+email_validator==2.1.1
+fastapi==0.111.0
+fastapi-cli==0.0.4
+ffmpy==0.3.2
+filelock==3.13.1
+fonttools==4.53.0
+fsspec==2024.2.0
+ftfy==6.2.0
+gradio==4.36.1
+gradio_client==1.0.1
+h11==0.14.0
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.4
+idna==3.7
+importlib_resources==6.4.0
+intel-openmp==2021.4.0
+Jinja2==3.1.3
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.0
+mdurl==0.1.2
+mkl==2021.4.0
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.3
+opencv-python==4.10.0.84
+orjson==3.10.5
+packaging==24.1
+pandas==2.2.2
+pillow==10.2.0
+pydantic==2.7.4
+pydantic_core==2.18.4
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.9
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+rich==13.7.1
+rpds-py==0.18.1
+ruff==0.4.9
+scipy==1.14.0
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+starlette==0.37.2
+sympy==1.12
+tbb==2021.11.0
+tiktoken==0.7.0
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.3.1
+torchaudio==2.3.1
+torchvision==0.18.1
+tqdm==4.66.4
+transformers==4.43.3
+typer==0.12.3
+typing_extensions==4.9.0
+tzdata==2024.1
+ujson==5.10.0
+urllib3==2.2.1
+uvicorn==0.30.1
+watchfiles==0.22.0
+websockets==11.0.3

rickyart.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0e2263b57fb66b48c9b09bee7b01b5fd8d708a6c52754265fd35052424d82ee
+size 3819

tony-diterlizzi-s-planescape-art.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:718333440c0986d401954d82acd2b2e0e8222f6b4d8587d4332c16bc9191cba4
+size 3819