From 512ffa90304a6bbfa023129c6f81761f22d702c8 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Fri, 9 Sep 2022 20:34:32 +0530
Subject: [PATCH] UI for setting the beta/main channel

---
 ui/index.html | 42 +++++++++++++++++++++++++++++++++++++++++-
 ui/server.py  | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/ui/index.html b/ui/index.html
index 03c3ce0c..5204a698 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -269,7 +269,7 @@
                     <div id="server-status-color">&nbsp;</div>
                     <span id="server-status-msg">Stable Diffusion is starting..</span>
                 </div>
-                <h1>Stable Diffusion UI <small>v2.08 (beta)</small></h1>
+                <h1>Stable Diffusion UI <small>v2.09</small></h1>
             </div>
             <div id="editor-inputs">
                 <div id="editor-inputs-prompt" class="row">
@@ -349,6 +349,8 @@
                     <li><input id="use_cpu" name="use_cpu" type="checkbox"> <label for="use_cpu">Use CPU instead of GPU (warning: this will be *very* slow)</label></li>
                     <li><input id="use_full_precision" name="use_full_precision" type="checkbox"> <label for="use_full_precision">Use full precision (for GPU-only. warning: this will consume more VRAM)</label></li>
                     <!-- <li><input id="allow_nsfw" name="allow_nsfw" type="checkbox"> <label for="allow_nsfw">Allow NSFW Content (You confirm you are above 18 years of age)</label></li> -->
+                    <br/>
+                    <li><input id="use_beta_channel" name="use_beta_channel" type="checkbox"> <label for="use_beta_channel">🔥Beta channel. Get the latest features immediately (but could be less stable). Please restart the program after changing this.</label></li>
                 </ul>
             </div>
 
@@ -416,6 +418,7 @@ let useFullPrecisionField = document.querySelector('#use_full_precision')
 let saveToDiskField = document.querySelector('#save_to_disk')
 let diskPathField = document.querySelector('#diskPath')
 // let allowNSFWField = document.querySelector("#allow_nsfw")
+let useBetaChannelField = document.querySelector("#use_beta_channel")
 let promptStrengthField = document.querySelector('#prompt_strength')
 let promptStrengthValueLabel = document.querySelector('#prompt_strength_value')
 
@@ -907,6 +910,42 @@ function updatePromptStrength() {
 promptStrengthField.addEventListener('input', updatePromptStrength)
 updatePromptStrength()
 
+useBetaChannelField.addEventListener('click', async function(e) {
+    let updateBranch = (this.checked ? 'beta' : 'main')
+
+    try {
+        let res = await fetch('/app_config', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json'
+            },
+            body: JSON.stringify({
+                'update_branch': updateBranch
+            })
+        })
+        res = await res.json()
+
+        console.log('set config status response', res)
+    } catch (e) {
+        console.log('set config status error', e)
+    }
+})
+
+async function getAppConfig() {
+    try {
+        let res = await fetch('/app_config')
+        config = await res.json()
+
+        if (config.update_branch === 'beta') {
+            useBetaChannelField.checked = true
+        }
+
+        console.log('get config status response', config)
+    } catch (e) {
+        console.log('get config status error', e)
+    }
+}
+
 function checkRandomSeed() {
     if (randomSeedField.checked) {
         seedField.disabled = true
@@ -1134,6 +1173,7 @@ async function loadModifiers() {
 async function init() {
     await loadModifiers()
     await getDiskPath()
+    await getAppConfig()
 
     setInterval(healthCheck, HEALTH_PING_INTERVAL * 1000)
     healthCheck()
diff --git a/ui/server.py b/ui/server.py
index 9a55cb5f..e287f5f1 100644
--- a/ui/server.py
+++ b/ui/server.py
@@ -1,3 +1,4 @@
+import json
 import traceback
 
 import sys
@@ -9,6 +10,8 @@ print('started in ', SCRIPT_DIR)
 SD_UI_DIR = os.getenv('SD_UI_PATH', None)
 sys.path.append(os.path.dirname(SD_UI_DIR))
 
+CONFIG_DIR = os.path.join(SD_UI_DIR, '..', 'scripts')
+
 OUTPUT_DIRNAME = "Stable Diffusion UI" # in the user's home folder
 
 from fastapi import FastAPI, HTTPException
@@ -43,6 +46,9 @@ class ImageRequest(BaseModel):
     use_cpu: bool = False
     use_full_precision: bool = False
 
+class SetAppConfigRequest(BaseModel):
+    update_branch: str = "main"
+
 @app.get('/')
 def read_root():
     return FileResponse(os.path.join(SD_UI_DIR, 'index.html'))
@@ -100,6 +106,48 @@ async def image(req : ImageRequest):
         print(traceback.format_exc())
         return HTTPException(status_code=500, detail=str(e))
 
+@app.post('/app_config')
+async def setAppConfig(req : SetAppConfigRequest):
+    try:
+        config = {
+            'update_branch': req.update_branch
+        }
+
+        config_json_str = json.dumps(config)
+        config_bat_str = f'@set update_branch={req.update_branch}'
+        config_sh_str = f'export update_branch={req.update_branch}'
+
+        config_json_path = os.path.join(CONFIG_DIR, 'config.json')
+        config_bat_path = os.path.join(CONFIG_DIR, 'config.bat')
+        config_sh_path = os.path.join(CONFIG_DIR, 'config.sh')
+
+        with open(config_json_path, 'w') as f:
+            f.write(config_json_str)
+
+        with open(config_bat_path, 'w') as f:
+            f.write(config_bat_str)
+
+        with open(config_sh_path, 'w') as f:
+            f.write(config_sh_str)
+
+        return {'OK'}
+    except Exception as e:
+        print(traceback.format_exc())
+        return HTTPException(status_code=500, detail=str(e))
+
+@app.get('/app_config')
+def getAppConfig():
+    try:
+        config_json_path = os.path.join(CONFIG_DIR, 'config.json')
+
+        with open(config_json_path, 'r') as f:
+            config_json_str = f.read()
+            config = json.loads(config_json_str)
+            return config
+    except Exception as e:
+        print(traceback.format_exc())
+        return HTTPException(status_code=500, detail=str(e))
+
 @app.get('/media/ding.mp3')
 def read_ding():
     return FileResponse(os.path.join(SD_UI_DIR, 'media/ding.mp3'))