Update Stable Diffusion 2.1 example (#2951)

ROCm · Apr 9, 2024 · ef52abd · ef52abd
1 parent 7374d52
commit ef52abd
Show file tree

Hide file tree

Showing 6 changed files with 495 additions and 135 deletions.
diff --git a/examples/diffusion/python_stable_diffusion_21/README.md b/examples/diffusion/python_stable_diffusion_21/README.md
@@ -1,6 +1,6 @@
 # Stable Diffusion 2.1
 
-This version was tested with [rocm 5.7](https://github.com/ROCmSoftwarePlatform/AMDMIGraphX/tree/rocm-5.7.0) revision.
+This version was tested with [rocm 6.0](https://github.com/ROCmSoftwarePlatform/AMDMIGraphX/tree/rocm-6.0.0) revision.
 
 ## Jupyter notebook
 
@@ -33,7 +33,7 @@ export PYTHONPATH=/opt/rocm/lib:$PYTHONPATH
 Get models with optimum
 
 ```bash
-optimum-cli export onnx --model stabilityai/stable-diffusion-2-1 models/sd21-onnx
+optimum-cli export onnx --model stabilityai/stable-diffusion-2-1 models/sd21-onnx --task stable-diffusion
 ```
 *Note: `models/sd21-onnx` will be used in the scripts.*
 
@@ -61,7 +61,7 @@ pip install -r gradio_requirements.txt
 Usage
 
 ```bash
-python gradio_app.py
+python gradio_app.py -p "a photograph of an astronaut riding a horse" --seed 13
 ```
 
 This will load the models (which can take several minutes), and when the setup is ready, starts a server on `http://127.0.0.1:7860`.
diff --git a/examples/diffusion/python_stable_diffusion_21/gradio_app.py b/examples/diffusion/python_stable_diffusion_21/gradio_app.py
@@ -1,7 +1,7 @@
 #####################################################################################
 # The MIT License (MIT)
 #
-# Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -22,13 +22,17 @@
 # THE SOFTWARE.
 #####################################################################################
 
-from txt2img import StableDiffusionMGX
+from txt2img import StableDiffusionMGX, get_args
 import gradio as gr
 
 
 def main():
+    args = get_args()
     # Note: This will load the models, which can take several minutes
-    sd = StableDiffusionMGX()
+    sd = StableDiffusionMGX(args.onnx_model_path, args.compiled_model_path,
+                            args.fp16, args.force_compile,
+                            args.exhaustive_tune)
+    sd.warmup(5)
 
     def gr_wrapper(prompt, negative_prompt, steps, seed, scale):
         result = sd.run(str(prompt), str(negative_prompt), int(steps),
@@ -38,12 +42,14 @@ def gr_wrapper(prompt, negative_prompt, steps, seed, scale):
     demo = gr.Interface(
         gr_wrapper,
         [
-            gr.Textbox(value="a photograph of an astronaut riding a horse",
-                       label="Prompt"),
-            gr.Textbox(value="", label="Negative prompt (Optional)"),
-            gr.Slider(1, 100, step=1, value=20, label="Number of steps"),
-            gr.Textbox(value=13, label="Random seed"),
-            gr.Slider(1, 20, step=0.1, value=7.0, label="Guidance scale"),
+            gr.Textbox(value=args.prompt, label="Prompt"),
+            gr.Textbox(value=args.negative_prompt,
+                       label="Negative prompt (Optional)"),
+            gr.Slider(
+                1, 100, step=1, value=args.steps, label="Number of steps"),
+            gr.Textbox(value=args.seed, label="Random seed"),
+            gr.Slider(
+                1, 20, step=0.1, value=args.scale, label="Guidance scale"),
         ],
         "image",
     )

diff --git a/...table_diffusion_21/gradio_reqirements.txt → ...able_diffusion_21/gradio_requirements.txt b/...table_diffusion_21/gradio_reqirements.txt → ...able_diffusion_21/gradio_requirements.txt
@@ -1,7 +1,7 @@
 #####################################################################################
 # The MIT License (MIT)
 #
-# Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -21,5 +21,5 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 #####################################################################################
--f requirements.txt
+-r requirements.txt
 gradio
diff --git a/examples/diffusion/python_stable_diffusion_21/requirements.txt b/examples/diffusion/python_stable_diffusion_21/requirements.txt
@@ -1,7 +1,7 @@
 #####################################################################################
 # The MIT License (MIT)
 #
-# Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -21,6 +21,8 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 #####################################################################################
+-f https://repo.radeon.com/rocm/manylinux/rocm-rel-6.0/
+torch==2.1.1
 accelerate
 diffusers
 optimum[onnxruntime]