first commit

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+checkpoints/BK-SDM-Small_iter50000
+checkpoints/BK-SDM-Small_iter50000.tar.gz

app.py

@@ -0,0 +1,64 @@
+import gradio as gr
+from pathlib import Path
+from demo import SdmCompressionDemo
+
+if __name__ == "__main__":
+    servicer = SdmCompressionDemo()
+    example_list = servicer.get_example_list()
+
+    with gr.Blocks(theme='nota-ai/theme') as demo:
+        gr.Markdown(Path('docs/header.md').read_text())
+        gr.Markdown(Path('docs/description.md').read_text())
+        with gr.Row():
+            with gr.Column(variant='panel',scale=30):
+
+                text = gr.Textbox(label="Input Prompt", max_lines=5, placeholder="Enter your prompt")
+                
+                with gr.Row().style(equal_height=True):
+                    generate_original_button = gr.Button(value="Generate with Original Model", variant="primary")
+                    generate_compressed_button = gr.Button(value="Generate with Compressed Model", variant="primary")
+
+                with gr.Accordion("Advanced Settings", open=False):
+                    negative = gr.Textbox(label=f'Negative Prompt', placeholder=f'Enter aspects to remove (e.g., {"low quality"})')                       
+                    with gr.Row():                        
+                        guidance_scale = gr.Slider(label="Guidance Scale", value=7.5, minimum=4, maximum=11, step=0.5)
+                        steps = gr.Slider(label="Denoising Steps", value=25, minimum=10, maximum=75, step=5)
+                        seed = gr.Slider(0, 999999, label='Random Seed', value=1234, step=1)
+
+                with gr.Tab("Example Prompts"):
+                    examples = gr.Examples(examples=example_list, inputs=[text])
+                    
+            with gr.Column(variant='panel',scale=35):
+                # Define original model output components
+                gr.Markdown('<h2 align="center">Original Stable Diffusion 1.4</h2>')
+                original_model_output = gr.Image(label="Original Model")
+                with gr.Row().style(equal_height=True):
+                    original_model_test_time = gr.Textbox(value="", label="Inference Time (sec)")
+                    original_model_error = gr.Markdown()
+
+            with gr.Column(variant='panel',scale=35):
+                # Define compressed model output components
+                gr.Markdown('<h2 align="center">Compressed Stable Diffusion (Ours)</h2>')
+                compressed_model_output = gr.Image(label="Compressed Model")                
+                with gr.Row().style(equal_height=True):
+                    compressed_model_test_time = gr.Textbox(value="", label="Inference Time (sec)")
+                    compressed_model_error = gr.Markdown()
+        
+        inputs = [text, negative, guidance_scale, steps, seed]  
+
+        # Click the generate button for original model 
+        original_model_outputs = [original_model_output, original_model_error, original_model_test_time]    
+        text.submit(servicer.infer_original_model, inputs=inputs, outputs=original_model_outputs)        
+        generate_original_button.click(servicer.infer_original_model, inputs=inputs, outputs=original_model_outputs)
+        
+        # Click the generate button for compressed model
+        compressed_model_outputs = [compressed_model_output, compressed_model_error, compressed_model_test_time]
+        text.submit(servicer.infer_compressed_model, inputs=inputs, outputs=compressed_model_outputs)    
+        generate_compressed_button.click(servicer.infer_compressed_model, inputs=inputs, outputs=compressed_model_outputs)
+
+        gr.Markdown(Path('docs/footer.md').read_text())
+
+    demo.queue(concurrency_count=1)
+    # demo.launch()
+    demo.launch(share=True, auth=("test", "testasdf@@19"))
+

demo.py

@@ -0,0 +1,100 @@
+from diffusers import StableDiffusionPipeline, UNet2DConditionModel
+import torch
+import copy
+
+import time
+
+ORIGINAL_CHECKPOINT_ID = "CompVis/stable-diffusion-v1-4"
+COMPRESSED_UNET_PATH = "checkpoints/BK-SDM-Small_iter50000"
+
+DEVICE='cuda'
+# DEVICE='cpu'
+
+class SdmCompressionDemo:
+    def __init__(self) -> None:
+        self.device = DEVICE
+        self.torch_dtype = torch.float16 if 'cuda' in self.device else torch.float32
+        
+        self.pipe_original = StableDiffusionPipeline.from_pretrained(ORIGINAL_CHECKPOINT_ID,
+                                                                     torch_dtype=self.torch_dtype)        
+        self.pipe_compressed = copy.deepcopy(self.pipe_original)
+        self.pipe_compressed.unet = UNet2DConditionModel.from_pretrained(COMPRESSED_UNET_PATH, 
+                                                                         subfolder="unet",
+                                                                         torch_dtype=self.torch_dtype)
+        if 'cuda' in self.device:
+            self.pipe_original = self.pipe_original.to(self.device)
+            self.pipe_compressed = self.pipe_compressed.to(self.device)
+        self.device_msg = 'Tested on GPU.' if 'cuda' in self.device else 'Tested on CPU.'
+
+    def generate_image(self, pipe, text, negative, guidance_scale, steps, seed):
+        generator = torch.Generator(self.device).manual_seed(seed)
+        start = time.time()
+        result = pipe(text, negative_prompt = negative, generator = generator, 
+                      guidance_scale = guidance_scale, num_inference_steps = steps)
+        test_time = time.time() - start   
+
+        image = result.images[0]
+        nsfw_detected = result.nsfw_content_detected[0]     
+        print(f"text {text} | Processed time: {test_time} sec | nsfw_flag {nsfw_detected}")
+        print(f"negative {negative} | guidance_scale {guidance_scale} | steps {steps} ")
+        print("===========")
+
+        return image, nsfw_detected, format(test_time, ".2f")
+    
+    def error_msg(self, nsfw_detected):
+        if nsfw_detected:
+            return self.device_msg+" Black images are returned when potential harmful content is detected. Try different prompts or seeds."
+        else:
+            return self.device_msg
+
+    def check_invalid_input(self, text):
+        if text == '':
+            return True
+
+    def infer_original_model(self, text, negative, guidance_scale, steps, seed):
+        print(f"=== ORIG model --- seed {seed}")        
+        if self.check_invalid_input(text):
+            return None, "Please enter the input prompt.", None
+        output_image, nsfw_detected, test_time = self.generate_image(self.pipe_original, 
+                                                       text, negative, guidance_scale, steps, seed)
+        
+        return output_image, self.error_msg(nsfw_detected), test_time
+        
+    def infer_compressed_model(self, text, negative, guidance_scale, steps, seed):
+        print(f"=== COMPRESSED model --- seed {seed}")
+        if self.check_invalid_input(text):
+            return None, "Please enter the input prompt.", None
+        output_image, nsfw_detected, test_time = self.generate_image(self.pipe_compressed, 
+                                                       text, negative, guidance_scale, steps, seed) 
+
+        return output_image, self.error_msg(nsfw_detected), test_time
+    
+
+    def get_example_list(self):
+        return [
+            'a tropical bird sitting on a branch of a tree',
+            'many decorative umbrellas hanging up',
+            'an orange cat staring off with pretty eyes',
+            'beautiful woman face with fancy makeup',
+            'a decorated living room with a stylish feel',
+            'a black vase holding a bouquet of roses',
+            'very elegant bedroom featuring natural wood',
+            'buffet-style food including cake and cheese',
+            'a tall castle sitting under a cloudy sky',
+            'closeup of a brown bear sitting in a grassy area',
+            'a large basket with many fresh vegetables',    
+            'house being built with lots of wood',
+            'a close up of a pizza with several toppings',   
+            'a golden vase with many different flows',    
+            'a statue of a lion face attached to brick wall',    
+            'something that looks particularly interesting',
+            'table filled with a variety of different dishes', 
+            'a cinematic view of a large snowy peak',
+            'a grand city in the year 2100, hyper realistic',
+            'a blue eyed baby girl looking at the camera',
+        ]
+
+
+
+
+                

docs/description.md

@@ -0,0 +1,6 @@
+This demo showcases a compressed Stable Diffusion model (SDM) for general-purpose text-to-image synthesis. Our lightest model (**BK-SDM-Small**) achieves **36% reduced** parameters and latency. This model is bulit with (i) removing several residual and attention blocks from the U-Net of SDM and (ii) distillation pretraining on only 0.22M LAION pairs (fewer than 0.1% of the full training set). Despite very limited training resources, our model can imitate the original SDM by benefiting from transferred knowledge.
+
+<!-- <center>
+    <img src="docs/fig_model.png" width="70%">
+</center> -->
+

docs/footer.md

@@ -0,0 +1,5 @@
+<p align="center">
+    <a href="https://netspresso.ai/"><img src="https://huggingface.co/spaces/nota-ai/theme/resolve/main/docs/logo/nota_favicon_800x800.png" width="96px" height="96px"></a>
+</p>
+
+<br/>

docs/header.md

@@ -0,0 +1,3 @@
+# <center>Lightweight Text-to-Image Generation Demo</center>
+
+

requirements.txt

@@ -0,0 +1,3 @@
+torch==1.13.1
+gradio==3.31.0
+diffusers==0.15.0.dev0