Upload 6 files

README.md

@@ -1,13 +1,13 @@
 ---
-title: Stable Diff Multilingual V0.1
-emoji: 😻
+title: SDXL Lightning
+emoji: ⚡
 colorFrom: yellow
-colorTo: indigo
+colorTo: gray
 sdk: gradio
-sdk_version: 4.20.1
+sdk_version: 4.19.1
 app_file: app.py
 pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,61 @@
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from PIL import Image
+
+from model import *
+
+# SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
+
+# Constants
+# base = "stabilityai/stable-diffusion-xl-base-1.0"
+# repo = "ByteDance/SDXL-Lightning"
+# checkpoints = {
+#     "1-Step" : ["sdxl_lightning_1step_unet_x0.safetensors", 1],
+#     "2-Step" : ["sdxl_lightning_2step_unet.safetensors", 2],
+#     "4-Step" : ["sdxl_lightning_4step_unet.safetensors", 4],
+#     "8-Step" : ["sdxl_lightning_8step_unet.safetensors", 8],
+# }
+# loaded = None
+
+
+# Ensure model and scheduler are initialized in GPU-enabled function
+# if torch.cuda.is_available():
+#     pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
+
+    
+# Function 
+# @spaces.GPU(enable_queue=True)
+
+def generate_image(prompt):
+    
+    return prompt_to_img(prompt)[0]
+
+
+
+# Gradio Interface
+description = """
+This demo utilizes the SDXL-Lightning model by ByteDance, which is a lightning-fast text-to-image generative model capable of producing high-quality images in 4 steps.
+As a community effort, this demo was put together by AngryPenguin. Link to model: https://huggingface.co/ByteDance/SDXL-Lightning
+"""
+
+with gr.Blocks(css="style.css") as demo:
+    gr.HTML("<h1><center>Text-to-Image with SDXL-Lightning ⚡</center></h1>")
+    gr.Markdown(description)
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Textbox(label='Enter your prompt (English)', scale=8)
+            ckpt = gr.Dropdown(label='Select inference steps',choices=['1-Step', '2-Step', '4-Step', '8-Step'], value='4-Step', interactive=True)
+            submit = gr.Button(scale=1, variant='primary')
+    img = gr.Image(label='SDXL-Lightning Generated Image')
+
+    prompt.submit(fn=generate_image,
+                 inputs=[prompt],
+                 outputs=img,
+                 )
+    submit.click(fn=generate_image,
+                 inputs=[prompt],
+                 outputs=img,
+                 )
+    
+demo.queue().launch()

model.py

@@ -0,0 +1,187 @@
+import os
+from PIL import Image, ImageDraw
+import cv2
+import numpy as np
+from IPython.display import HTML
+from base64 import b64encode
+
+import torch
+from torch import autocast
+from torch.nn import functional as F
+from diffusers import StableDiffusionPipeline, AutoencoderKL
+from diffusers import UNet2DConditionModel, PNDMScheduler, LMSDiscreteScheduler
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+#from transformers import CLIPTextModel, CLIPTokenizer
+from tqdm.auto import tqdm
+from huggingface_hub import notebook_login
+
+import weights
+
+device = 'cpu'
+
+from Multilingual_CLIP.multilingual_clip import Config_MCLIP
+import transformers
+import torch
+
+
+class MultilingualCLIP(transformers.PreTrainedModel):
+    config_class = Config_MCLIP.MCLIPConfig
+
+    def __init__(self, config, *args, **kwargs):
+        super().__init__(config, *args, **kwargs)
+        self.transformer = transformers.AutoModel.from_pretrained(config.modelBase)
+        self.LinearTransformation = torch.nn.Linear(in_features=config.transformerDimensions,
+                                                    out_features=config.numDims)
+
+    def forward(self, txt, tokenizer, device):
+        txt_tok = tokenizer(txt, padding='max_length', max_length=77, truncation=True, return_tensors='pt').to(device)
+        embs = self.transformer(**txt_tok)
+        embs = embs[0]
+        att = txt_tok['attention_mask']
+        embs = (embs * att.unsqueeze(2)) / att.sum(dim=1)[:, None].unsqueeze(2)
+        return self.LinearTransformation(embs)
+
+    @classmethod
+    def _load_state_dict_into_model(cls, model, state_dict, pretrained_model_name_or_path, _fast_init=True):
+        model.load_state_dict(state_dict)
+        return model, [], [], []
+
+
+import torch
+import torch.nn as nn
+
+# Define the adaptation layer, 'checkpoint_9.pth'
+class AdaptationLayer(nn.Module):
+  def __init__(self, input_dim, output_dim):
+    super(AdaptationLayer, self).__init__()
+    self.fc1 = nn.Linear(input_dim, output_dim*2)
+    torch.nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity='relu')
+    self.bn1 = nn.BatchNorm1d(77)
+
+    self.fc2 = nn.Linear(input_dim*2, output_dim*2)
+    torch.nn.init.kaiming_uniform_(self.fc2.weight, nonlinearity='relu')
+    self.bn2 = nn.BatchNorm1d(77)
+
+    self.fc3 = nn.Linear(input_dim*2, output_dim)
+    torch.nn.init.kaiming_uniform_(self.fc3.weight, nonlinearity='relu')
+    self.bn3 = nn.BatchNorm1d(77)
+
+    self.fc4 = nn.Linear(input_dim, output_dim)
+    torch.nn.init.kaiming_uniform_(self.fc4.weight, nonlinearity='relu')
+    self.bn4 = nn.BatchNorm1d(77)
+
+    self.fc5 = nn.Linear(input_dim, output_dim)
+
+  def forward(self, x):
+    x = nn.functional.normalize(x, p=2.0, dim=1, eps=1e-12, out=None)
+    x = torch.relu(self.bn1(self.fc1(x)))
+    x = torch.relu(self.bn2(self.fc2(x)))
+    x = torch.relu(self.bn3(self.fc3(x)))
+    x = torch.relu(self.bn4(self.fc4(x)))
+
+    return self.fc5(x)
+  
+
+adapt_model  = AdaptationLayer(768,768)
+adapt_model.to(device)
+state_dict = torch.load('weights/checkpoint_9.pth')
+adapt_model.load_state_dict(state_dict)
+
+# 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', use_auth_token=True)
+vae = vae.to(device)
+
+# 2. Load the tokenizer and text encoder to tokenize and encode the text.
+tokenizer = text_tokenizer
+text_encoder = text_model
+
+# 3. The UNet model for generating the latents.
+unet = UNet2DConditionModel.from_pretrained(
+    'CompVis/stable-diffusion-v1-4', subfolder='unet', use_auth_token=True)
+unet = unet.to(device)
+
+# 4. Create a scheduler for inference
+scheduler = LMSDiscreteScheduler(
+    beta_start=0.00085, beta_end=0.012,
+    beta_schedule='scaled_linear', num_train_timesteps=1000)
+
+
+def get_text_embeds(prompt):
+  with torch.no_grad():
+    text_embeddings = text_model(prompt, text_tokenizer, device)
+    text_embeddings = adapt_model(text_embeddings)
+
+  # Do the same for unconditional embeddings
+  with torch.no_grad():
+    uncond_embeddings = text_model([''] * len(prompt), text_tokenizer, device)
+    uncond_embeddings = adapt_model(uncond_embeddings)
+
+  # Cat for final embeddings
+  text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+  return text_embeddings
+
+
+def produce_latents(text_embeddings, height=512, width=512,
+                    num_inference_steps=50, guidance_scale=7.5, latents=None):
+  if latents is None:
+    latents = torch.randn((text_embeddings.shape[0] // 2, unet.in_channels, \
+                           height // 8, width // 8))
+  latents = latents.to(device)
+
+  scheduler.set_timesteps(num_inference_steps)
+  latents = latents * scheduler.sigmas[0]
+
+  with autocast('cpu'):
+    for i, t in tqdm(enumerate(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 = latent_model_input / ((sigma**2 + 1) ** 0.5)
+
+      # predict the noise residual
+      with torch.no_grad():
+        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings.to(device))['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, i, latents)['prev_sample']
+
+  return latents
+
+
+def decode_img_latents(latents):
+  latents = 1 / 0.18215 * latents
+
+  with torch.no_grad():
+    imgs = vae.decode(latents)
+
+  imgs = (imgs / 2 + 0.5).clamp(0, 1)
+  imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
+  imgs = (imgs * 255).round().astype('uint8')
+  pil_images = [Image.fromarray(image) for image in imgs]
+  return pil_images
+
+def prompt_to_img(prompts, height=512, width=512, num_inference_steps=50,
+                  guidance_scale=7.5, latents=None):
+  if isinstance(prompts, str):
+    prompts = [prompts]
+
+  # Prompts -> text embeds
+  text_embeds = get_text_embeds(prompts)
+
+  # Text embeds -> img latents
+  latents = produce_latents(
+      text_embeds, height=height, width=width, latents=latents,
+      num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)
+
+  # Img latents -> imgs
+  imgs = decode_img_latents(latents)
+  
+  return imgs
+
+
+

requirements.txt

@@ -0,0 +1,5 @@
+transformers
+diffusers
+torch
+accelerate
+gradio

safety_checker.py

@@ -0,0 +1,137 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel
+
+
+def cosine_distance(image_embeds, text_embeds):
+    normalized_image_embeds = nn.functional.normalize(image_embeds)
+    normalized_text_embeds = nn.functional.normalize(text_embeds)
+    return torch.mm(normalized_image_embeds, normalized_text_embeds.t())
+
+
+class StableDiffusionSafetyChecker(PreTrainedModel):
+    config_class = CLIPConfig
+
+    _no_split_modules = ["CLIPEncoderLayer"]
+
+    def __init__(self, config: CLIPConfig):
+        super().__init__(config)
+
+        self.vision_model = CLIPVisionModel(config.vision_config)
+        self.visual_projection = nn.Linear(
+            config.vision_config.hidden_size, config.projection_dim, bias=False
+        )
+
+        self.concept_embeds = nn.Parameter(
+            torch.ones(17, config.projection_dim), requires_grad=False
+        )
+        self.special_care_embeds = nn.Parameter(
+            torch.ones(3, config.projection_dim), requires_grad=False
+        )
+
+        self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
+        self.special_care_embeds_weights = nn.Parameter(
+            torch.ones(3), requires_grad=False
+        )
+
+    @torch.no_grad()
+    def forward(self, clip_input, images):
+        pooled_output = self.vision_model(clip_input)[1]  # pooled_output
+        image_embeds = self.visual_projection(pooled_output)
+
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        special_cos_dist = (
+            cosine_distance(image_embeds, self.special_care_embeds)
+            .cpu()
+            .float()
+            .numpy()
+        )
+        cos_dist = (
+            cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy()
+        )
+
+        result = []
+        batch_size = image_embeds.shape[0]
+        for i in range(batch_size):
+            result_img = {
+                "special_scores": {},
+                "special_care": [],
+                "concept_scores": {},
+                "bad_concepts": [],
+            }
+
+            # increase this value to create a stronger `nfsw` filter
+            # at the cost of increasing the possibility of filtering benign images
+            adjustment = 0.0
+
+            for concept_idx in range(len(special_cos_dist[0])):
+                concept_cos = special_cos_dist[i][concept_idx]
+                concept_threshold = self.special_care_embeds_weights[concept_idx].item()
+                result_img["special_scores"][concept_idx] = round(
+                    concept_cos - concept_threshold + adjustment, 3
+                )
+                if result_img["special_scores"][concept_idx] > 0:
+                    result_img["special_care"].append(
+                        {concept_idx, result_img["special_scores"][concept_idx]}
+                    )
+                    adjustment = 0.01
+
+            for concept_idx in range(len(cos_dist[0])):
+                concept_cos = cos_dist[i][concept_idx]
+                concept_threshold = self.concept_embeds_weights[concept_idx].item()
+                result_img["concept_scores"][concept_idx] = round(
+                    concept_cos - concept_threshold + adjustment, 3
+                )
+                if result_img["concept_scores"][concept_idx] > 0:
+                    result_img["bad_concepts"].append(concept_idx)
+
+            result.append(result_img)
+
+        has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result]
+
+        return has_nsfw_concepts
+
+    @torch.no_grad()
+    def forward_onnx(self, clip_input: torch.FloatTensor, images: torch.FloatTensor):
+        pooled_output = self.vision_model(clip_input)[1]  # pooled_output
+        image_embeds = self.visual_projection(pooled_output)
+
+        special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds)
+        cos_dist = cosine_distance(image_embeds, self.concept_embeds)
+
+        # increase this value to create a stronger `nsfw` filter
+        # at the cost of increasing the possibility of filtering benign images
+        adjustment = 0.0
+
+        special_scores = (
+            special_cos_dist - self.special_care_embeds_weights + adjustment
+        )
+        # special_scores = special_scores.round(decimals=3)
+        special_care = torch.any(special_scores > 0, dim=1)
+        special_adjustment = special_care * 0.01
+        special_adjustment = special_adjustment.unsqueeze(1).expand(
+            -1, cos_dist.shape[1]
+        )
+
+        concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment
+        # concept_scores = concept_scores.round(decimals=3)
+        has_nsfw_concepts = torch.any(concept_scores > 0, dim=1)
+
+        images[has_nsfw_concepts] = 0.0  # black image
+
+        return images, has_nsfw_concepts

style.css

@@ -0,0 +1,12 @@
+.gradio-container {
+  max-width: 690px! important;
+}
+
+#share-btn-container{padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; max-width: 13rem; margin-left: auto;margin-top: 0.35em;}
+div#share-btn-container > div {flex-direction: row;background: black;align-items: center}
+#share-btn-container:hover {background-color: #060606}
+#share-btn {all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.5rem !important; padding-bottom: 0.5rem !important;right:0;font-size: 15px;}
+#share-btn * {all: unset}
+#share-btn-container div:nth-child(-n+2){width: auto !important;min-height: 0px !important;}
+#share-btn-container .wrap {display: none !important}
+#share-btn-container.hidden {display: none!important}