first

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -0,0 +1,127 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionXLPipeline
+from diffusers.schedulers import TCDScheduler
+import spaces
+from PIL import Image
+
+SAFETY_CHECKER = True
+
+# Constants
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+checkpoints = {
+    "2-Step": ["pcm_sdxl_smallcfg_2step_converted.safetensors", 2, 0.0],
+    "4-Step": ["pcm_sdxl_smallcfg_4step_converted.safetensors", 4, 0.0],
+    "8-Step": ["pcm_sdxl_smallcfg_8step_converted.safetensors", 8, 0.0],
+    "16-Step": ["pcm_sdxl_smallcfg_16step_converted.safetensors", 16, 0.0],
+    "Normal CFG 4-Step": ["pcm_sdxl_normalcfg_4step_converted.safetensors", 4, 7.5],
+    "Normal CFG 8-Step": ["pcm_sdxl_normalcfg_8step_converted.safetensors", 8, 7.5],
+    "Normal CFG 16-Step": ["pcm_sdxl_normalcfg_16step_converted.safetensors", 16, 7.5],
+    "LCM-Like LoRA": ["pcm_sdxl_lcmlike_lora_converted.safetensors", 16, 0.0],
+}
+
+
+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")
+
+if SAFETY_CHECKER:
+    from safety_checker import StableDiffusionSafetyChecker
+    from transformers import CLIPFeatureExtractor
+
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    ).to("cuda")
+    feature_extractor = CLIPFeatureExtractor.from_pretrained(
+        "openai/clip-vit-base-patch32"
+    )
+
+    def check_nsfw_images(
+        images: list[Image.Image],
+    ) -> tuple[list[Image.Image], list[bool]]:
+        safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
+        has_nsfw_concepts = safety_checker(
+            images=[images], clip_input=safety_checker_input.pixel_values.to("cuda")
+        )
+
+        return images, has_nsfw_concepts
+
+
+# Function
+@spaces.GPU(enable_queue=True)
+def generate_image(prompt, ckpt):
+    global loaded
+    print(prompt, ckpt)
+
+    checkpoint = checkpoints[ckpt][0]
+    num_inference_steps = checkpoints[ckpt][1]
+    guidance_scale = checkpoints[ckpt][2]
+
+    if loaded != num_inference_steps:
+        pipe.scheduler = TCDScheduler(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            timestep_spacing="trailing",
+        )
+        pipe.load_lora_weights(
+            "wangfuyun/PCM_Weights", weight_name=checkpoint, subfolder="sdxl"
+        )
+
+        loaded = num_inference_steps
+
+    results = pipe(
+        prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
+    )
+
+    if SAFETY_CHECKER:
+        images, has_nsfw_concepts = check_nsfw_images(results.images)
+        if any(has_nsfw_concepts):
+            gr.Warning("NSFW content detected.")
+            return Image.new("RGB", (512, 512))
+        return images[0]
+    return results.images[0]
+
+
+# Gradio Interface
+
+css = """
+.gradio-container {
+  max-width: 60rem !important;
+}
+"""
+with gr.Blocks(css=css) as demo:
+    gr.HTML("<h1><center>SDXL-Lightning ⚡</center></h1>")
+    gr.HTML(
+        "<p><center>Lightning-fast text-to-image generation</center></p><p><center><a href='https://huggingface.co/ByteDance/SDXL-Lightning'>https://huggingface.co/ByteDance/SDXL-Lightning</a></center></p>"
+    )
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Textbox(label="Enter your prompt (English)", scale=8)
+            ckpt = gr.Dropdown(
+                label="Select inference steps",
+                choices=list(checkpoints.keys()),
+                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, ckpt],
+        outputs=img,
+    )
+    submit.click(
+        fn=generate_image,
+        inputs=[prompt, ckpt],
+        outputs=img,
+    )
+
+demo.queue().launch()

requirements.txt

@@ -0,0 +1,8 @@
+diffusers==0.28.0
+datasets
+transformers
+accelerate
+peft
+xformers
+gradio==4.32.2
+spaces

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