add KDPM2Discrete and a force flag for setup_venv (#1044)

* add KDPM2Discrete and a force flag for setup_venv

* add KDPM2Discrete and a force flag for setup_venv
also made sure that Python 3.11 is used for the venv as 3.10
doesn't work anymore

* add KDPM2Discrete and a force flag for setup_venv
also made sure that Python 3.11 is used for the venv as 3.10
doesn't work anymore

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_txt2img.py

@@ -9,6 +9,7 @@ from diffusers import (
     DDIMScheduler,
     PNDMScheduler,
     LMSDiscreteScheduler,
+    KDPM2DiscreteScheduler,
     EulerDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
@@ -30,6 +31,7 @@ class Text2ImagePipeline(StableDiffusionPipeline):
             DDIMScheduler,
             PNDMScheduler,
             LMSDiscreteScheduler,
+            KDPM2DiscreteScheduler,
             EulerDiscreteScheduler,
             EulerAncestralDiscreteScheduler,
             DPMSolverMultistepScheduler,

apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_utils.py

@@ -9,6 +9,7 @@ from diffusers import (
     DDIMScheduler,
     PNDMScheduler,
     LMSDiscreteScheduler,
+    KDPM2DiscreteScheduler,
     EulerDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
@@ -40,6 +41,7 @@ class StableDiffusionPipeline:
             DDIMScheduler,
             PNDMScheduler,
             LMSDiscreteScheduler,
+            KDPM2DiscreteScheduler,
             EulerDiscreteScheduler,
             EulerAncestralDiscreteScheduler,
             DPMSolverMultistepScheduler,
@@ -182,6 +184,7 @@ class StableDiffusionPipeline:
             DDIMScheduler,
             PNDMScheduler,
             LMSDiscreteScheduler,
+            KDPM2DiscreteScheduler,
             EulerDiscreteScheduler,
             EulerAncestralDiscreteScheduler,
             DPMSolverMultistepScheduler,

apps/stable_diffusion/src/schedulers/sd_schedulers.py

@@ -3,6 +3,7 @@ from diffusers import (
     PNDMScheduler,
     DDIMScheduler,
     DPMSolverMultistepScheduler,
+    KDPM2DiscreteScheduler,
     EulerDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
 )
@@ -17,6 +18,10 @@ def get_schedulers(model_id):
         model_id,
         subfolder="scheduler",
     )
+    schedulers["KDPM2Discrete"] = KDPM2DiscreteScheduler.from_pretrained(
+        model_id,
+        subfolder="scheduler",
+    )
     schedulers["LMSDiscrete"] = LMSDiscreteScheduler.from_pretrained(
         model_id,
         subfolder="scheduler",

apps/stable_diffusion/web/ui/txt2img_ui.py

@@ -93,6 +93,7 @@ with gr.Blocks(title="Text-to-Image") as txt2img_web:
                                 "DDIM",
                                 "PNDM",
                                 "LMSDiscrete",
+                                "KDPM2Discrete",
                                 "DPMSolverMultistep",
                                 "EulerDiscrete",
                                 "EulerAncestralDiscrete",

setup_venv.ps1

@@ -1,9 +1,62 @@
+<#
+.SYNOPSIS
+  A script to update and install the SHARK runtime and its dependencies.
+
+.DESCRIPTION
+  This script updates and installs the SHARK runtime and its dependencies.
+  It checks the Python version installed and installs any required build
+  dependencies into a Python virtual environment.
+  If that environment does not exist, it creates it.
+  
+.PARAMETER update-src
+  updates to latest version from git .\source
+
+.PARAMETER force
+  removes and recreates venv to force update all dependencies
+  
+.EXAMPLE
+  .\setup_venv.ps1 --force
+
+.EXAMPLE
+  .\setup_venv.ps1 --update-src
+
+.INPUTS
+  None
+
+.OUTPUTS
+  None
+
+.NOTES
+  Version        1.0
+  Author         powderluv, xzuyn
+  Creation Date  2023-02-17
+  PurposeChange Initial script development
+
+#>
+
 param([string]$arguments)
 
 if ($arguments -eq "--update-src"){
 	git pull
 }
 
+if ($arguments -eq "--force"){
+	if (Test-Path env:VIRTUAL_ENV) {
+        Write-Host "deactivating..."
+        Deactivate
+    }
+    
+    if (Test-Path .\shark.venv\) {
+        Write-Host "removing and recreating venv..."
+        Remove-Item .\shark.venv -Force -Recurse
+        if (Test-Path .\shark.venv\) {
+            Write-Host 'could not remove .\shark-venv - please try running ".\setup_venv.ps1 --force" again!'
+            break
+        }
+    }
+}
+
+
 #Write-Host "Installing python"
 
 #Start-Process winget install Python.Python.3.10 '/quiet InstallAllUsers=1 PrependPath=1' -wait -NoNewWindow
@@ -31,11 +84,17 @@ else
 
 Write-Host "Python version found is"
 Write-Host $p
-
+if ($p -notlike "*3.11*")
+{
+    Write-Host "Please install Python 3.11 and try again"
+    break
+}
 
 Write-Host "Installing Build Dependencies"
-python -m venv .\shark.venv\
+py -3.11 -m venv .\shark.venv\
 .\shark.venv\Scripts\activate
+python -m pip install --upgrade pip
+pip install wheel
 pip install -r requirements.txt
 pip install --pre torch-mlir torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu -f https://llvm.github.io/torch-mlir/package-index/
 pip install --upgrade -f https://nod-ai.github.io/SHARK-Runtime/pip-release-links.html iree-compiler iree-runtime