Merge branch 'huggingface-spaces' into main

.gitignore

@@ -1,3 +1,6 @@
 __pycache__/
-checkpoints/BK-SDM-Small_iter50000
+checkpoints/BK-SDM-Small_iter50000/unet/*
 checkpoints/BK-SDM-Small_iter50000.tar.gz
+app.sh
+
+!checkpoints/BK-SDM-Small_iter50000/unet/.gitkeep

Dockerfile

@@ -0,0 +1,7 @@
+FROM nvcr.io/nvidia/pytorch:22.03-py3
+
+ARG DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+RUN apt-get install ffmpeg libsm6 libxext6 tmux git -y
+
+WORKDIR /workspace

app.py

@@ -1,7 +1,27 @@
-import gradio as gr
+import os
+import subprocess
 from pathlib import Path
+
+import gradio as gr
+
 from demo import SdmCompressionDemo
 
+dest_path_config = Path('checkpoints/BK-SDM-Small_iter50000/unet/config.json')
+dest_path_torch_ckpt = Path('checkpoints/BK-SDM-Small_iter50000/unet/diffusion_pytorch_model.bin')
+BK_SDM_CONFIG_URL: str = os.getenv('BK_SDM_CONFIG_URL', None)
+BK_SDM_TORCH_CKPT_URL: str = os.getenv('BK_SDM_TORCH_CKPT_URL', None)
+assert BK_SDM_CONFIG_URL is not None
+assert BK_SDM_TORCH_CKPT_URL is not None
+
+subprocess.call(
+    f"wget --no-check-certificate -O {dest_path_config} {BK_SDM_CONFIG_URL}",
+    shell=True
+)
+subprocess.call(
+    f"wget --no-check-certificate -O {dest_path_torch_ckpt} {BK_SDM_TORCH_CKPT_URL}",
+    shell=True
+)
+
 if __name__ == "__main__":
     servicer = SdmCompressionDemo()
     example_list = servicer.get_example_list()
@@ -10,25 +30,25 @@ if __name__ == "__main__":
         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):
+            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():                        
+                    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):
+
+            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")
@@ -36,29 +56,28 @@ if __name__ == "__main__":
                     original_model_test_time = gr.Textbox(value="", label="Inference Time (sec)")
                     original_model_error = gr.Markdown()
 
-            with gr.Column(variant='panel',scale=35):
+            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")                
+                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)        
+        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)    
+        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.launch(server_name="0.0.0.0", server_port=7861)

docker-compose.yml

@@ -0,0 +1,16 @@
+version: "3.9"
+
+# docker compose run --service-ports --name bk-sdm-hksong bk-sdm bash
+
+services:
+  bk-sdm:
+    image: bk-sdm:dev
+    build: ./
+    container_name: bk-sdm-hksong
+    ipc: host
+    ports:
+      - "7861:7861"
+    volumes:
+      - /data2/hksong/bk-sdm:/workspace
+      - /data2/hksong/DATA:/DATA
+      - /data2/hksong/LOG:/LOG

requirements.txt

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