custom model integration done

Author: divamgupta

backends/stable_diffusion_tf/stable_diffusion_tf/constants.py

@@ -1755,6 +1755,7 @@
     49407,
     49407,
 ]
+
 _ALPHAS_CUMPROD = [
     0.99915,
     0.998296,

electron_app/src/App.vue

@@ -76,6 +76,10 @@
                 <History :app_state="app_state"></History>
             </template>
 
+            <template v-slot:settings>
+                <Settings :app_state="app_state"></Settings>
+            </template>
+
             <template v-slot:logs>
                 <div class="animatable_content_box ">
                     <p>Logs : </p>
@@ -111,6 +115,8 @@ import Outpainting from './components/Outpainting.vue'
 import UpscaleImage from './components/UpscaleImage.vue'
 import Inpainting from "./components/Inpainting.vue"
 import History from './components/History.vue'
+import Settings from './components/Settings.vue'
+
 
 import LoaderModal from './components_bare/LoaderModal.vue'
 import Vue from "vue"
@@ -131,7 +137,8 @@ export default
         Img2Img,
         Outpainting,
         UpscaleImage, 
-        Inpainting
+        Inpainting,
+        Settings
     },
 
     mounted() {
@@ -157,8 +164,19 @@ export default
 
 
         let data = window.ipcRenderer.sendSync('load_data');
-        if( data.history)
-            Vue.set(this.app_state , 'history' , data.history)
+        if(!data.history){
+            data.history = {}
+        }
+        if(!data.settings){
+            data.settings = {}
+        }
+        if(!data.custom_models){
+            data.custom_models = {}
+        }
+        if( data ){
+            Vue.set(this.app_state , 'app_data' , data)
+        }
+            
      
     },
 
@@ -183,10 +201,10 @@ export default
             deep: true
         } , 
 
-        'app_state.history': {
+        'app_state.app_data': {
 
             handler: function(new_value) {
-                window.ipcRenderer.sendSync('save_data', {"history": new_value });
+                window.ipcRenderer.sendSync('save_data', new_value );
             },
             deep: true
         } , 
@@ -293,7 +311,7 @@ export default
             show_splash_screen : true , // is showing the loading splash screen
             logs : "",
             global_loader_modal_msg : "",
-            history : {},
+            app_data: {history : {}},
         };
 
         return {

electron_app/src/StableDiffusion.vue

@@ -131,8 +131,7 @@ export default {
 
             if(prompt_params.negative_prompt){
                 prompt_params.negative_prompt = remove_non_ascii(prompt_params.negative_prompt)
-            }
-                
+            }                
 
             this.last_iter_t = Date.now()
             this.generated_by = generated_by;

electron_app/src/assets/css/theme.css

@@ -12,9 +12,9 @@
 
 @media (prefers-color-scheme: light) {
      :root {
-        --slider-progress: #D0A3FD;
+        --slider-progress: #8cb8f2;
         --slider-progress_end: #F1F4FA;
-        --slider-circle: #9131F1;
+        --slider-circle: #3174f1;
         --slider-border: white;
     }
     body {
@@ -136,9 +136,9 @@
 
 @media (prefers-color-scheme: dark) {
      :root {
-        --slider-progress: #5304a3;
+        --slider-progress: #0444a3;
         --slider-progress_end: #303030;
-        --slider-circle: #3a076d;
+        --slider-circle: #07296d;
         --slider-border: #222222;
     }
     body {

electron_app/src/components/History.vue

@@ -4,7 +4,7 @@
         <div @click="toggle_order()" style="float:right; margin-bottom: 20px;" class="l_button">
           {{this.app_state.show_history_in_oldest_first ? "Oldest": "Newest"}} First
         </div>
-        <div v-if="Object.values(app_state.history).length > 0">
+        <div v-if="Object.values(app_state.app_data.history).length > 0">
             <div v-for="history_box in get_history()" :key="history_box.key" style="clear: both;">
 
                 <div @click="delete_hist(history_box.key)" style="float:right; margin-top: 10px;"  class="l_button">Delete</div>
@@ -73,11 +73,11 @@ export default {
     },
     methods: {
         delete_hist(k){
-            Vue.delete( this.app_state.history , k );
+            Vue.delete( this.app_state.app_data.history , k );
         },
 
         get_history() {
-          let history = Object.values(this.app_state.history);
+          let history = Object.values(this.app_state.app_data.history);
           return this.app_state.show_history_in_oldest_first ? history : history.reverse();
         },
 

electron_app/src/components/Img2Img.vue

@@ -138,6 +138,7 @@ export default {
 
             is_negative_prompt_avail : false, 
             negative_prompt : "",
+            selected_model : 'Default'
         };
     },
     methods: {
@@ -180,8 +181,15 @@ export default {
                 batch_size : this.batch_size , 
                 img_strength : this.inp_img_strength,
                 input_image : input_image,
-                is_inpaint : this.is_inpaint,
+                is_inpaint : this.is_inpaint                
             }
+
+            if(this.selected_model && this.selected_model != "Default" && this.app_state.app_data.custom_models[this.selected_model] ){
+                params.model_id = -1;
+                params.custom_model_path =  this.app_state.app_data.custom_models[this.selected_model].path;
+            }
+
+            
             if(this.is_inpaint)
                 params['mask_image'] = mask_img;
 
@@ -199,7 +207,6 @@ export default {
                 on_img(img_path){
                     that.generated_images.push(img_path);
 
-                
                     let p = {
                             "prompt":that.prompt , "seed": seed, "key":history_key , "imgs" : [] , "inp_img": input_image_with_mask,
                             "dif_steps" : that.dif_steps , "inp_img_strength" : that.inp_img_strength, "model_version": that.stable_diffusion.model_version , "guidence_scale" : that.guidence_scale , 
@@ -209,12 +216,12 @@ export default {
                     if(that.is_negative_prompt_avail)
                         p['negative_prompt'] = that.negative_prompt;
 
-                    if(!(that.app_state.history[history_key]))
-                        Vue.set(that.app_state.history, history_key , p );
+                    if(!(that.app_state.app_data.history[history_key]))
+                        Vue.set(that.app_state.app_data.history, history_key , p );
                     
-                    that.app_state.history[history_key].imgs.push(img_path)
+                    that.app_state.app_data.history[history_key].imgs.push(img_path)
 
-                    console.log(that.app_state.history)
+                    console.log(that.app_state.app_data.history)
 
                 },
                 on_progress(p){

electron_app/src/components/ImgGenerate.vue

@@ -72,7 +72,7 @@
 
             <div v-if="generated_images.length == 1" >
                 <center>
-                    <ImageItem :app_state="app_state"  :path="generated_images[0]" :style_obj="{ 'width': 'calc(100vh - 380px )' , 'margin-top': '60px' }"></ImageItem>
+                    <ImageItem :app_state="app_state"  :path="generated_images[0]" :style_obj="{ 'width': 'calc(100vh - 390px )' , 'margin-top': '60px' }"></ImageItem>
                 </center>
             </div>
 
@@ -152,6 +152,7 @@ export default {
             modifiers : require("../modifiers.json"),
             is_negative_prompt_avail : false, 
             negative_prompt : "",
+            selected_model : 'Default'
         };
         
     },
@@ -179,8 +180,12 @@ export default {
                 scale : this.guidence_scale , 
                 ddim_steps : this.dif_steps, 
                 num_imgs : this.num_imgs , 
-                batch_size : this.batch_size , 
+                batch_size : this.batch_size 
+            }
 
+            if(this.selected_model && this.selected_model != "Default" && this.app_state.app_data.custom_models[this.selected_model] ){
+                params.model_id = -1;
+                params.custom_model_path =  this.app_state.app_data.custom_models[this.selected_model].path;
             }
 
             if(this.is_negative_prompt_avail)
@@ -201,23 +206,23 @@ export default {
                 on_img(img_path){
                     that.generated_images.push(img_path);
 
-                    if(!(that.app_state.history[history_key])){
+                    if(!(that.app_state.app_data.history[history_key])){
                         let p = {
                             "prompt":that.prompt , "seed": seed, "img_w":that.img_w , "img_h":that.img_h ,  "key":history_key , "imgs" : [],
-                            "guidence_scale" : that.guidence_scale , "dif_steps" : that.dif_steps  
+                            "guidence_scale" : that.guidence_scale , "dif_steps" : that.dif_steps 
                         }
                         if(that.stable_diffusion.model_version)
                             p['model_version'] = that.stable_diffusion.model_version;
                         if(that.is_negative_prompt_avail)
                             p['negative_prompt'] = that.negative_prompt;
-                        Vue.set(that.app_state.history, history_key , p);
+                        Vue.set(that.app_state.app_data.history, history_key , p);
                     }
                         
 
                     
-                    that.app_state.history[history_key].imgs.push(img_path)
+                    that.app_state.app_data.history[history_key].imgs.push(img_path)
 
-                    console.log(that.app_state.history)
+                    console.log(that.app_state.app_data.history)
 
                 },
                 on_progress(p ){

electron_app/src/components/Settings.vue

@@ -0,0 +1,94 @@
+<template>
+    <div  class="animatable_content_box ">
+
+        <div  v-if="!is_custom_model_loading">
+            <h1>Settings</h1>
+            <br>
+            <h2>General Settings</h2>
+
+            <br><hr>
+            <div class="l_button button_colored" style="float:right" @click="add_model"  >Add New Model</div>
+            <h2>Custom Models</h2>
+            
+            <!-- <br> -->
+            <hr> 
+            <div v-for="model in Object.values(this.app_state.app_data.custom_models)" :key="model.name">
+                <div class="l_button" @click="delete_model(model.name)" style="float:right">Remove</div>
+                <p> Name : {{model.name}} </p>
+                <p> Path : {{model.orig_path}} </p>
+                <hr> 
+            </div>
+        </div>
+        <div v-else>
+            <LoaderModal :loading_percentage="-1" loading_title="Importing model"></LoaderModal>
+        </div>
+
+    </div>
+</template>
+<script>
+import LoaderModal from '../components_bare/LoaderModal.vue'
+import Vue from 'vue'
+
+
+export default {
+    name: 'Settings',
+    props: {
+        app_state : Object , 
+    },
+    components: {LoaderModal},
+    mounted() {
+
+    },
+    data() {
+        return {
+            is_custom_model_loading:false,
+        };
+    },
+    methods: {
+        delete_model(k){
+            let model_path = this.app_state.app_data.custom_models[k].path;
+            window.ipcRenderer.sendSync('delete_file',  model_path );
+            Vue.delete( this.app_state.app_data.custom_models , k );
+        },
+        add_model(){
+            let that = this;
+            
+            let pytorch_model_path = window.ipcRenderer.sendSync('file_dialog',  "ckpt_file" );
+            if(!pytorch_model_path)
+                return;
+
+            if(pytorch_model_path == "NULL")
+                return;
+
+            let model_name = pytorch_model_path.replaceAll("\\" , "/").split("/");
+            model_name = model_name[model_name.length - 1].split(".")[0]
+
+
+            if(model_name.trim() == ""){
+                alert("Put non empty model name");
+                return;
+            }
+
+            if(that.app_state.app_data.custom_models[model_name]){
+                alert("this model name "+ model_name +" already exists")
+                return;
+            }
+
+            that.is_custom_model_loading = true
+            window.ipcRenderer.invoke('add_custom_pytorch_models', pytorch_model_path , model_name ).then((result) => {
+                that.is_custom_model_loading = false
+                if(result.success){
+                    Vue.set( that.app_state.app_data.custom_models , model_name ,  {"name":model_name , "path": result.model_path, "orig_path" : pytorch_model_path }  );
+                } else {
+                    alert("Error " + result.error )
+                }
+                
+            })
+        }
+    },
+}
+</script>
+<style>
+</style>
+<style scoped>
+</style>

electron_app/src/components_bare/ApplicationFrame.vue

@@ -10,6 +10,7 @@
                                 <font-awesome-icon icon="bars" />
                             </div>
                         </template>
+                        
                         <b-dropdown-item-button @click="$emit('menu_item_click_help',{})">Help</b-dropdown-item-button>
                         <b-dropdown-item-button @click="$emit('menu_item_click_discord',{})">Start Discord Chat</b-dropdown-item-button>
                         <b-dropdown-item-button @click="selectTab('logs')">Show Logs</b-dropdown-item-button>
@@ -19,6 +20,8 @@
 
                         <b-dropdown-item-button @click="$emit('menu_item_click_about',{})">About</b-dropdown-item-button>
 
+                        <b-dropdown-item-button @click="selectTab('settings')">Settings</b-dropdown-item-button>
+
                         <b-dropdown-item-button @click="$emit('menu_item_click_close',{})">Close</b-dropdown-item-button>
                         <!-- #TODO set these menu items via python -->
 
@@ -65,13 +68,19 @@
                     <slot name="outpainting"></slot>
                 </div>
 
+                <div style="display:none"  :class="{ bl_display : selected_tab === 'settings'  }" >
+                    <slot name="settings"></slot>
+                </div>
+
+
 
 
                 <div  v-if="selected_tab === 'logs' " style="display:none"  :class="{ bl_display : selected_tab === 'logs'  }" >
                     <slot name="logs"></slot>
                 </div>
 
 
+