inference test

app.py

@@ -5,6 +5,9 @@ from utils import call
 from diffusers.pipelines import StableDiffusionXLPipeline
 StableDiffusionXLPipeline.__call__ = call
 import os
+from trainscripts.textsliders.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV
+
+
 os.environ['CURL_CA_BUNDLE'] = ''
 model_map = {'Age' : 'models/age.pt', 
              'Chubby': 'models/chubby.pt',
@@ -36,7 +39,7 @@ class Demo:
         self.generating = False
         self.device = 'cuda'
         self.weight_dtype = torch.float16
-        self.pipe = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=self.weight_dtype)
+        self.pipe = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=self.weight_dtype).to(self.device)
 
         with gr.Blocks() as demo:
             self.layout()
@@ -57,7 +60,7 @@ class Demo:
 
                 with gr.Row():
 
-                    self.explain_infr = gr.Markdown(value='This is a demo of [Concept Sliders: LoRA Adaptors for Precise Control in Diffusion Models](https://sliders.baulab.info/). To try out a model that can control a particular concept, select a model and enter any prompt.  For example, if you select the model "Surprised Look" you can generate images for the prompt "A picture of a person, realistic, 8k" and compare the slider effect to the image generated by original model.  We have also provided several other pre-fine-tuned models like "repair" sliders to repair flaws in SDXL generated images (Check out the "Pretrained Sliders" drop-down). You can also train and run your own custom sliders. Check out the "train" section for custom concept slider training.')
+                    self.explain_infr = gr.Markdown(value='This is a demo of [Concept Sliders: LoRA Adaptors for Precise Control in Diffusion Models](https://sliders.baulab.info/). To try out a model that can control a particular concept, select a model and enter any prompt, choose a seed, and finally choose the SDEdit timestep for structural preservation. Higher SDEdit timesteps results in more structural change. For example, if you select the model "Surprised Look" you can generate images for the prompt "A picture of a person, realistic, 8k" and compare the slider effect to the image generated by original model.  We have also provided several other pre-fine-tuned models like "repair" sliders to repair flaws in SDXL generated images (Check out the "Pretrained Sliders" drop-down). You can also train and run your own custom sliders. Check out the "train" section for custom concept slider training.')
 
                 with gr.Row():
 
@@ -82,6 +85,20 @@ class Demo:
                                 label="Seed",
                                 value=12345
                             )
+                            
+                            self.slider_scale_infr = gr.Number(
+                                label="Slider Scale",
+                                value=2,
+                                info="Larger slider scale result in stronger edit"
+                            )
+
+                            
+                            self.start_noise_infr = gr.Slider(
+                                600, 900, 
+                                value=750, 
+                                label="SDEdit Timestep", 
+                                info="Choose smaller values for more structural preservation"
+                            )
 
                     with gr.Column(scale=2):
 
@@ -162,6 +179,8 @@ class Demo:
         self.infr_button.click(self.inference, inputs = [
             self.prompt_input_infr,
             self.seed_infr,
+            self.start_noise_infr,
+            self.slider_scale_infr,
             self.model_dropdown
             ],
             outputs=[
@@ -217,7 +236,7 @@ class Demo:
 #         return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom model in the "Test" tab'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
         return None
 
-    def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
+    def inference(self, prompt, seed, start_noise, scale, model_name, pbar = gr.Progress(track_tqdm=True)):
         
         seed = seed or 12345
 
@@ -225,41 +244,53 @@ class Demo:
 
         model_path = model_map[model_name]
         
-        checkpoint = torch.load(model_path)
-        
-        return None
-
-#         finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()
-
-#         torch.cuda.empty_cache()
+        unet = self.pipe.unet
+        network_type = "c3lier"
+        if 'full' in model_path:
+            train_method = 'full'
+        elif 'noxattn' in model_path:
+            train_method = 'noxattn'
+        elif 'xattn' in model_path:
+            train_method = 'xattn'
+            network_type = 'lierla'
+        else:
+            train_method = 'noxattn'
+
+        modules = DEFAULT_TARGET_REPLACE
+        if network_type == "c3lier":
+            modules += UNET_TARGET_REPLACE_MODULE_CONV
+
+        name = os.path.basename(model_path)
+        rank = 4
+        alpha = 1
+        if 'rank4' in model_path:
+            rank = 4
+        if 'rank8' in model_path:
+            rank = 8
+        if 'alpha1' in model_path:
+            alpha = 1.0
+        network = LoRANetwork(
+                unet,
+                rank=rank,
+                multiplier=1.0,
+                alpha=alpha,
+                train_method=train_method,
+            ).to(self.device, dtype=self.weight_dtype)
+        network.load_state_dict(torch.load(model_path))
 
-#         images = self.diffuser(
-#             prompt,
-#             n_steps=50,
-#             generator=generator
-#         )
 
+        generator = torch.manual_seed(seed)
+        edited_image = pipe(prompt, num_images_per_prompt=1, num_inference_steps=50, generator=generator, network=network, start_noise=start_noise, scale=scale, unet=unet).images[0]
         
-#         orig_image = images[0][0]
-
-#         torch.cuda.empty_cache()
-
-#         generator = torch.manual_seed(seed)
-
-#         with finetuner:
-
-#             images = self.diffuser(
-#                 prompt,
-#                 n_steps=50,
-#                 generator=generator
-#             )
-
-#         edited_image = images[0][0]
-
-#         del finetuner
-#         torch.cuda.empty_cache()
-
-#         return edited_image, orig_image
-
+        generator = torch.manual_seed(seed)
+        original_image = pipe(prompt, num_images_per_prompt=1, num_inference_steps=50, generator=generator, network=network, start_noise=start_noise, scale=0, unet=unet).images[0]
+        
+        del unet, network
+        unet = None
+        network = None
+        pipe = None
+        torch.cuda.empty_cache()
+        
+        return edited_image, original_image 
 
 demo = Demo()