adding refiner option

app.py

@@ -16,13 +16,15 @@ def models_change(model, scheduler, config):
     config = set_config(config, 'model', model)
 
     use_safetensors = False
+    refiner = ""
     
     # no model selected (because this is UI init run)
     if type(model) != list and str(model) != 'None':
         
         use_safetensors = str(models[model]['use_safetensors'])
         model_description = models[model]['description']
-        
+        refiner = models[model]['refiner']
+                
         # if no scheduler is selected, choose the default one for this model
         if scheduler == None:
             
@@ -34,11 +36,12 @@ def models_change(model, scheduler, config):
         
     config["use_safetensors"] = str(use_safetensors)
     config["scheduler"] = str(scheduler)
+    config["refiner"] = str(refiner)
     
     # safety_checker_change(in_safety_checker.value, config)
     # requires_safety_checker_change(in_requires_safety_checker.value, config)
 
-    return model_description, use_safetensors, scheduler, config, str(config), assemble_code(config)
+    return model_description, refiner, use_safetensors, scheduler, config, str(config), assemble_code(config)
 
 def data_type_change(data_type, config):
 
@@ -132,6 +135,15 @@ def run_inference(config, config_history, progress=gr.Progress(track_tqdm=True))
             torch_dtype = get_data_type(config["data_type"]), 
             variant = get_variant(config["variant"])).to(config["device"])
         
+        if config['refiner'] != '':
+            refiner = DiffusionPipeline.from_pretrained(
+                    config['refiner'],
+                    text_encoder_2=pipeline.text_encoder_2,
+                    vae=pipeline.vae,
+                    torch_dtype=get_data_type(config["data_type"]),
+                    use_safetensors=get_bool(config["use_safetensors"]), 
+                    variant = get_variant(config["variant"])).to(config["device"])
+            
         if str(config["safety_checker"]).lower() == 'false':
             pipeline.safety_checker = None 
 
@@ -151,11 +163,18 @@ def run_inference(config, config_history, progress=gr.Progress(track_tqdm=True))
             negative_prompt = config["negative_prompt"],
             generator = generator,
             num_inference_steps = int(config["inference_steps"]), 
-            guidance_scale = float(config["guidance_scale"])).images[0]
-        
+            guidance_scale = float(config["guidance_scale"])).images
+
+        if config['refiner'] != '':
+            image = refiner(
+                prompt = config["prompt"],
+                num_inference_steps = int(config["inference_steps"]), 
+                image=image,
+            ).images
+            
         config_history.append(config.copy())
 
-        return image, dict_list_to_markdown_table(config_history), config_history
+        return image[0], dict_list_to_markdown_table(config_history), config_history
     
     else:
     
@@ -191,6 +210,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
     with gr.Row():
         with gr.Column(scale=1):
             in_use_safetensors = gr.Radio(label="Use safe tensors:", choices=["True", "False"], interactive=False)
+            in_model_refiner = gr.Textbox(value="", label="Refiner")
         with gr.Column(scale=1):
             in_safety_checker = gr.Radio(label="Enable safety checker:", value=config.value["safety_checker"], choices=["True", "False"])
             in_requires_safety_checker = gr.Radio(label="Requires safety checker:", value=config.value["requires_safety_checker"], choices=["True", "False"])
@@ -229,7 +249,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
     in_data_type.change(data_type_change, inputs=[in_data_type, config], outputs=[config, out_config, out_code])
     in_allow_tensorfloat32.change(tensorfloat32_change, inputs=[in_allow_tensorfloat32, config], outputs=[config, out_config, out_code])
     in_variant.change(variant_change, inputs=[in_variant, config], outputs=[config, out_config, out_code])
-    in_models.change(models_change, inputs=[in_models, in_schedulers, config], outputs=[out_model_description, in_use_safetensors, in_schedulers, config, out_config, out_code])
+    in_models.change(models_change, inputs=[in_models, in_schedulers, config], outputs=[out_model_description, in_model_refiner, in_use_safetensors, in_schedulers, config, out_config, out_code])
     in_safety_checker.change(safety_checker_change, inputs=[in_safety_checker, config], outputs=[config, out_config, out_code])
     in_requires_safety_checker.change(requires_safety_checker_change, inputs=[in_requires_safety_checker, config], outputs=[config, out_config, out_code])
     in_schedulers.change(schedulers_change, inputs=[in_schedulers, config], outputs=[out_scheduler_description, config, out_config, out_code])
@@ -249,6 +269,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
             in_devices,
             in_use_safetensors,
             in_data_type,
+            in_model_refiner,
             in_variant,
             in_safety_checker,
             in_requires_safety_checker,

appConfig.json

@@ -4,47 +4,56 @@
         "sd-dreambooth-library/solo-levelling-art-style": {
             "use_safetensors": false,
             "description": "see https://huggingface.co/sd-dreambooth-library/solo-levelling-art-style",
-            "scheduler": "DDIMScheduler"
+            "scheduler": "DDIMScheduler",
+            "refiner": ""
         },
         "CompVis/stable-diffusion-v1-4": {
             "use_safetensors": true,
             "description": "see https://huggingface.co/CompVis/stable-diffusion-v1-4",
-            "scheduler": "EulerDiscreteScheduler"
+            "scheduler": "EulerDiscreteScheduler",
+            "refiner": ""
         },
         "runwayml/stable-diffusion-v1-5": {
             "use_safetensors": true,
             "description": "see https://huggingface.co/runwayml/stable-diffusion-v1-5",
-            "scheduler": "DDPMScheduler"
+            "scheduler": "DDPMScheduler",
+            "refiner": ""
         },
         "stabilityai/stable-diffusion-2-1": {
             "use_safetensors": true,
             "description": "see https://huggingface.co/stabilityai/stable-diffusion-2-1",
-            "scheduler": "DPMSolverMultistepScheduler"
+            "scheduler": "DPMSolverMultistepScheduler",
+            "refiner": ""
         },
         "stabilityai/stable-diffusion-xl-base-1.0": {
             "use_safetensors": true,
             "description": "see https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0",
-            "scheduler": "DDPMScheduler"
+            "scheduler": "DDPMScheduler",
+            "refiner": "stabilityai/stable-diffusion-xl-refiner-1.0"
         },
         "sd-dreambooth-library/house-emblem": {
             "use_safetensors": false,
             "description": "see https://huggingface.co/sd-dreambooth-library/house-emblem",
-            "scheduler": "DDPMScheduler"
+            "scheduler": "DDPMScheduler",
+            "refiner": ""
         },
         "Envvi/Inkpunk-Diffusion": {
-            "use_safetensors": true,
-            "description": "Another wonderful model",
-            "scheduler": "DDPMScheduler"
+            "use_safetensors": false,
+            "description": "see https://huggingface.co/Envvi/Inkpunk-Diffusion",
+            "scheduler": "DDPMScheduler",
+            "refiner": ""
         },
         "Stelath/textual_inversion_comic_strip_fp16": {
             "use_safetensors": true,
-            "description": "Another wonderful model",
-            "scheduler": "DDPMScheduler"
+            "description": "see https://huggingface.co/Stelath/textual_inversion_comic_strip_fp16",
+            "scheduler": "DDPMScheduler",
+            "refiner": ""
         },
         "sd-dreambooth-library/herge-style": {
-            "use_safetensors": true,
-            "description": "Another wonderful model",
-            "scheduler": "DDPMScheduler"
+            "use_safetensors": false,
+            "description": "see https://huggingface.co/sd-dreambooth-library/herge-style",
+            "scheduler": "DDPMScheduler",
+            "refiner": ""
         }
 
     },

config.py

@@ -42,6 +42,7 @@ def get_initial_config():
         "allow_tensorfloat32": allow_tensorfloat32,
         "use_safetensors": "False",
         "data_type": data_type,
+        "refiner": "",
         "safety_checker": "False",
         "requires_safety_checker": "False",
         "manual_seed": 42,
@@ -78,6 +79,7 @@ def get_config_from_url(initial_config, request: Request):
             return_config['device'],
             return_config['use_safetensors'],
             return_config['data_type'],
+            return_config['refiner'],
             return_config['variant'],
             return_config['safety_checker'],
             return_config['requires_safety_checker'],
@@ -139,12 +141,20 @@ def assemble_code(str_config):
             torch_dtype=data_type, 
             variant=variant).to(device)'''
     
+    if config['refiner'] != '':
+        code['051_refiner'] = f'''refiner = DiffusionPipeline.from_pretrained(
+                "{config['refiner']}",
+                text_encoder_2 = base.text_encoder_2,
+                vae = base.vae,
+                torch_dtype = data_type,
+                use_safetensors = use_safetensors,
+                variant=variant,
+            ).to(device)'''
+        
     code['054_requires_safety_checker'] = f'pipeline.requires_safety_checker = {config["requires_safety_checker"]}'
 
     if str(config["safety_checker"]).lower() == 'false':
         code['055_safety_checker'] = f'pipeline.safety_checker = None'
-    else:
-        code['055_safety_checker'] = ''
     
     code['060_scheduler'] = f'pipeline.scheduler = {config["scheduler"]}.from_config(pipeline.scheduler.config)'
     
@@ -165,6 +175,16 @@ def assemble_code(str_config):
         negative_prompt=negative_prompt,
         generator=generator,
         num_inference_steps=inference_steps, 
-        guidance_scale=guidance_scale).images[0]'''
-            
+        guidance_scale=guidance_scale).images
+        '''
+
+    if config['refiner'] != '':
+        code["110_run_refiner"] = f'''image = refiner(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=inference_steps
+            ).images[0]'''
+        
+    code["200_show_image"] = 'image[0]'
+    
     return '\r\n'.join(value[1] for value in sorted(code.items()))