fully working sd webui api on ray serve

fully working sd webui api on ray serve

Author: webcoderz

extensions-builtin/Lora/networks.py

@@ -15,6 +15,7 @@
 from typing import Union
 
 from modules import shared, devices, sd_models, errors, scripts, sd_hijack
+from modules.shared import shared_instance
 
 module_types = [
     network_lora.ModuleTypeLora(),
@@ -112,8 +113,10 @@ def match(match_list, regex_text):
 def assign_network_names_to_compvis_modules(sd_model):
     network_layer_mapping = {}
 
-    if shared.sd_model.is_sdxl:
-        for i, embedder in enumerate(shared.sd_model.conditioner.embedders):
+    #if shared.sd_model.is_sdxl:
+    if shared_instance.sd_model.is_sdxl:
+        #for i, embedder in enumerate(shared.sd_model.conditioner.embedders):
+        for i, embedder in enumerate(shared_instance.sd_model.conditioner.embedders):
             if not hasattr(embedder, 'wrapped'):
                 continue
 
@@ -122,12 +125,14 @@ def assign_network_names_to_compvis_modules(sd_model):
                 network_layer_mapping[network_name] = module
                 module.network_layer_name = network_name
     else:
-        for name, module in shared.sd_model.cond_stage_model.wrapped.named_modules():
+        #for name, module in shared.sd_model.cond_stage_model.wrapped.named_modules():
+        for name, module in shared_instance.sd_model.cond_stage_model.wrapped.named_modules():
             network_name = name.replace(".", "_")
             network_layer_mapping[network_name] = module
             module.network_layer_name = network_name
 
-    for name, module in shared.sd_model.model.named_modules():
+    #for name, module in shared.sd_model.model.named_modules():
+    for name, module in shared_instance.sd_model.model.named_modules():
         network_name = name.replace(".", "_")
         network_layer_mapping[network_name] = module
         module.network_layer_name = network_name
@@ -142,37 +147,46 @@ def load_network(name, network_on_disk):
     sd = sd_models.read_state_dict(network_on_disk.filename)
 
     # this should not be needed but is here as an emergency fix for an unknown error people are experiencing in 1.2.0
-    if not hasattr(shared.sd_model, 'network_layer_mapping'):
-        assign_network_names_to_compvis_modules(shared.sd_model)
+    #if not hasattr(shared.sd_model, 'network_layer_mapping'):
+        #assign_network_names_to_compvis_modules(shared.sd_model)
+    if not hasattr(shared_instance.sd_model, 'network_layer_mapping'):
+        assign_network_names_to_compvis_modules(shared_instance.sd_model)
 
     keys_failed_to_match = {}
-    is_sd2 = 'model_transformer_resblocks' in shared.sd_model.network_layer_mapping
+    #is_sd2 = 'model_transformer_resblocks' in shared.sd_model.network_layer_mapping
+    is_sd2 = 'model_transformer_resblocks' in shared_instance.sd_model.network_layer_mapping
 
     matched_networks = {}
 
     for key_network, weight in sd.items():
         key_network_without_network_parts, network_part = key_network.split(".", 1)
 
         key = convert_diffusers_name_to_compvis(key_network_without_network_parts, is_sd2)
-        sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+        #sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+        sd_module = shared_instance.sd_model.network_layer_mapping.get(key, None)
 
         if sd_module is None:
             m = re_x_proj.match(key)
             if m:
-                sd_module = shared.sd_model.network_layer_mapping.get(m.group(1), None)
+                #sd_module = shared.sd_model.network_layer_mapping.get(m.group(1), None)
+                sd_module = shared_instance.sd_model.network_layer_mapping.get(m.group(1), None)
+
 
         # SDXL loras seem to already have correct compvis keys, so only need to replace "lora_unet" with "diffusion_model"
         if sd_module is None and "lora_unet" in key_network_without_network_parts:
             key = key_network_without_network_parts.replace("lora_unet", "diffusion_model")
-            sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+            #sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+            sd_module = shared_instance.sd_model.network_layer_mapping.get(key, None)
         elif sd_module is None and "lora_te1_text_model" in key_network_without_network_parts:
             key = key_network_without_network_parts.replace("lora_te1_text_model", "0_transformer_text_model")
-            sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+            #sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+            sd_module = shared_instance.sd_model.network_layer_mapping.get(key, None)
 
             # some SD1 Loras also have correct compvis keys
             if sd_module is None:
                 key = key_network_without_network_parts.replace("lora_te1_text_model", "transformer_text_model")
-                sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+                #sd_module = shared.sd_model.network_layer_mapping.get(key, None)
+                sd_module = shared_instance.sd_model.network_layer_mapping.get(key, None)
 
         if sd_module is None:
             keys_failed_to_match[key_network] = key

extensions-builtin/Lora/ui_extra_networks_lora.py

@@ -6,7 +6,7 @@
 from modules import shared, ui_extra_networks
 from modules.ui_extra_networks import quote_js
 from ui_edit_user_metadata import LoraUserMetadataEditor
-
+from modules.shared import shared_instance
 
 class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
     def __init__(self):
@@ -52,17 +52,26 @@ def create_item(self, name, index=None, enable_filter=True):
 
         if shared.opts.lora_show_all or not enable_filter:
             pass
+#        elif sd_version == network.SdVersion.Unknown:
+#            model_version = network.SdVersion.SDXL if shared.sd_model.is_sdxl else network.SdVersion.SD2 if shared.sd_model.is_sd2 else network.SdVersion.SD1
+#            if model_version.name in shared.opts.lora_hide_unknown_for_versions:
+#                return None
+#        elif shared.sd_model.is_sdxl and sd_version != network.SdVersion.SDXL:
+#            return None
+#        elif shared.sd_model.is_sd2 and sd_version != network.SdVersion.SD2:
+#            return None
+#        elif shared.sd_model.is_sd1 and sd_version != network.SdVersion.SD1:
+#            return None
         elif sd_version == network.SdVersion.Unknown:
-            model_version = network.SdVersion.SDXL if shared.sd_model.is_sdxl else network.SdVersion.SD2 if shared.sd_model.is_sd2 else network.SdVersion.SD1
+            model_version = network.SdVersion.SDXL if shared_instance.sd_model.is_sdxl else network.SdVersion.SD2 if shared_instance.sd_model.is_sd2 else network.SdVersion.SD1
             if model_version.name in shared.opts.lora_hide_unknown_for_versions:
                 return None
-        elif shared.sd_model.is_sdxl and sd_version != network.SdVersion.SDXL:
+        elif shared_instance.sd_model.is_sdxl and sd_version != network.SdVersion.SDXL:
             return None
-        elif shared.sd_model.is_sd2 and sd_version != network.SdVersion.SD2:
+        elif shared_instance.sd_model.is_sd2 and sd_version != network.SdVersion.SD2:
             return None
-        elif shared.sd_model.is_sd1 and sd_version != network.SdVersion.SD1:
+        elif shared_instance.sd_model.is_sd1 and sd_version != network.SdVersion.SD1:
             return None
-
         return item
 
     def list_items(self):

modules/api/api.py

@@ -33,7 +33,7 @@
 import piexif
 import piexif.helper
 from contextlib import closing
-
+from modules.shared import shared_instance
 
 def script_name_to_index(name, scripts):
     try:
@@ -364,7 +364,8 @@ def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI):
         args.pop('save_images', None)
 
         with self.queue_lock:
-            with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
+            with closing(StableDiffusionProcessingTxt2Img(sd_model=shared_instance.sd_model, **args)) as p:
+            #with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
                 p.is_api = True
                 p.scripts = script_runner
                 p.outpath_grids = opts.outdir_txt2img_grids
@@ -424,7 +425,8 @@ def img2imgapi(self, img2imgreq: models.StableDiffusionImg2ImgProcessingAPI):
         args.pop('save_images', None)
 
         with self.queue_lock:
-            with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
+            #with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
+            with closing(StableDiffusionProcessingImg2Img(sd_model=shared_instance.sd_model, **args)) as p:
                 p.init_images = [decode_base64_to_image(x) for x in init_images]
                 p.is_api = True
                 p.scripts = script_runner
@@ -724,8 +726,10 @@ def train_hypernetwork(self, args: dict):
             except Exception as e:
                 error = e
             finally:
-                shared.sd_model.cond_stage_model.to(devices.device)
-                shared.sd_model.first_stage_model.to(devices.device)
+                #shared.sd_model.cond_stage_model.to(devices.device)
+                #shared.sd_model.first_stage_model.to(devices.device)
+                shared_instance.sd_model.cond_stage_model.to(devices.device)
+                shared_instance.sd_model.first_stage_model.to(devices.device)
                 if not apply_optimizations:
                     sd_hijack.apply_optimizations()
                 shared.state.end()

modules/api/raypi.py

@@ -40,7 +40,7 @@
 from modules import initialize_util
 from modules import script_callbacks
 import os
-
+from modules.shared import shared_instance
 import launch
 from ray import serve
 
@@ -349,8 +349,8 @@ def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI):
         send_images = args.pop('send_images', True)
         args.pop('save_images', None)
 
-        
-        with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
+        with closing(StableDiffusionProcessingTxt2Img(sd_model=shared_instance.sd_model, **args)) as p:          
+        #with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
             p.is_api = True
             p.scripts = script_runner
             p.outpath_grids = opts.outdir_txt2img_grids
@@ -409,8 +409,8 @@ def img2imgapi(self, img2imgreq: models.StableDiffusionImg2ImgProcessingAPI):
         send_images = args.pop('send_images', True)
         args.pop('save_images', None)
 
-        
-        with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
+        with closing(StableDiffusionProcessingImg2Img(sd_model=shared_instance.sd_model, **args)) as p:
+        #with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
             p.init_images = [decode_base64_to_image(x) for x in init_images]
             p.is_api = True
             p.scripts = script_runner
@@ -737,8 +737,10 @@ def train_hypernetwork(self, args: dict):
             except Exception as e:
                 error = e
             finally:
-                shared.sd_model.cond_stage_model.to(devices.device)
-                shared.sd_model.first_stage_model.to(devices.device)
+                #shared.sd_model.cond_stage_model.to(devices.device)
+                #shared.sd_model.first_stage_model.to(devices.device)
+                shared_instance.sd_model.cond_stage_model.to(devices.device)
+                shared_instance.sd_model.first_stage_model.to(devices.device)
                 if not apply_optimizations:
                     sd_hijack.apply_optimizations()
                 shared.state.end()

modules/hypernetworks/hypernetwork.py

@@ -18,7 +18,7 @@
 
 from collections import deque
 from statistics import stdev, mean
-
+from modules.shared import shared_instance
 
 optimizer_dict = {optim_name : cls_obj for optim_name, cls_obj in inspect.getmembers(torch.optim, inspect.isclass) if optim_name != "Optimizer"}
 
@@ -525,7 +525,8 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
 
     pin_memory = shared.opts.pin_memory
 
-    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)
+    #ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)
+    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared_instance.sd_model, cond_model=shared_instance.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)
 
     if shared.opts.save_training_settings_to_txt:
         saved_params = dict(
@@ -542,8 +543,10 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
 
     if unload:
         shared.parallel_processing_allowed = False
-        shared.sd_model.cond_stage_model.to(devices.cpu)
-        shared.sd_model.first_stage_model.to(devices.cpu)
+        #shared.sd_model.cond_stage_model.to(devices.cpu)
+        #shared.sd_model.first_stage_model.to(devices.cpu)
+        shared_instance.sd_model.cond_stage_model.to(devices.cpu)
+        shared_instance.sd_model.first_stage_model.to(devices.cpu)
 
     weights = hypernetwork.weights()
     hypernetwork.train()
@@ -614,16 +617,21 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
                     if use_weight:
                         w = batch.weight.to(devices.device, non_blocking=pin_memory)
                     if tag_drop_out != 0 or shuffle_tags:
-                        shared.sd_model.cond_stage_model.to(devices.device)
-                        c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)
-                        shared.sd_model.cond_stage_model.to(devices.cpu)
+                        #shared.sd_model.cond_stage_model.to(devices.device)
+                        shared_instance.sd_model.cond_stage_model.to(devices.device)
+                        c = shared_instance.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)
+                        shared_instance.sd_model.cond_stage_model.to(devices.cpu)
+                        #c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)
+                        #shared.sd_model.cond_stage_model.to(devices.cpu)
                     else:
                         c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory)
                     if use_weight:
-                        loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step
+                        loss = shared_instance.sd_model.weighted_forward(x, c, w)[0] / gradient_step
+                        #loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step
                         del w
                     else:
-                        loss = shared.sd_model.forward(x, c)[0] / gradient_step
+                        #loss = shared.sd_model.forward(x, c)[0] / gradient_step
+                        loss = shared_instance.sd_model.forward(x, c)[0] / gradient_step
                     del x
                     del c
 
@@ -683,11 +691,14 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
                     cuda_rng_state = None
                     if torch.cuda.is_available():
                         cuda_rng_state = torch.cuda.get_rng_state_all()
-                    shared.sd_model.cond_stage_model.to(devices.device)
-                    shared.sd_model.first_stage_model.to(devices.device)
+                    #shared.sd_model.cond_stage_model.to(devices.device)
+                    #shared.sd_model.first_stage_model.to(devices.device)
+                    shared_instance.sd_model.cond_stage_model.to(devices.device)
+                    shared_instance.sd_model.first_stage_model.to(devices.device)
 
                     p = processing.StableDiffusionProcessingTxt2Img(
-                        sd_model=shared.sd_model,
+                        #sd_model=shared.sd_model,
+                        sd_model=shared_instance.sd_model,
                         do_not_save_grid=True,
                         do_not_save_samples=True,
                     )
@@ -716,8 +727,10 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
                         image = processed.images[0] if len(processed.images) > 0 else None
 
                     if unload:
-                        shared.sd_model.cond_stage_model.to(devices.cpu)
-                        shared.sd_model.first_stage_model.to(devices.cpu)
+                        #shared.sd_model.cond_stage_model.to(devices.cpu)
+                        #shared.sd_model.first_stage_model.to(devices.cpu)
+                        shared_instance.sd_model.cond_stage_model.to(devices.cpu)
+                        shared_instance.sd_model.first_stage_model.to(devices.cpu)
                     torch.set_rng_state(rng_state)
                     if torch.cuda.is_available():
                         torch.cuda.set_rng_state_all(cuda_rng_state)
@@ -760,8 +773,10 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
 
     del optimizer
     hypernetwork.optimizer_state_dict = None  # dereference it after saving, to save memory.
-    shared.sd_model.cond_stage_model.to(devices.device)
-    shared.sd_model.first_stage_model.to(devices.device)
+    #shared.sd_model.cond_stage_model.to(devices.device)
+    #shared.sd_model.first_stage_model.to(devices.device)
+    shared_instance.sd_model.cond_stage_model.to(devices.device)
+    shared_instance.sd_model.first_stage_model.to(devices.device)
     shared.parallel_processing_allowed = old_parallel_processing_allowed
 
     return hypernetwork, filename

modules/hypernetworks/ui.py

@@ -3,7 +3,7 @@
 import gradio as gr
 import modules.hypernetworks.hypernetwork
 from modules import devices, sd_hijack, shared
-
+from modules.shared import shared_instance
 not_available = ["hardswish", "multiheadattention"]
 keys = [x for x in modules.hypernetworks.hypernetwork.HypernetworkModule.activation_dict if x not in not_available]
 
@@ -32,7 +32,9 @@ def train_hypernetwork(*args):
     except Exception:
         raise
     finally:
-        shared.sd_model.cond_stage_model.to(devices.device)
-        shared.sd_model.first_stage_model.to(devices.device)
+        #shared.sd_model.cond_stage_model.to(devices.device)
+        #shared.sd_model.first_stage_model.to(devices.device)
+        shared_instance.sd_model.cond_stage_model.to(devices.device)
+        shared_instance.sd_model.first_stage_model.to(devices.device)
         sd_hijack.apply_optimizations()