multiple models dropdown

app.py

@@ -67,13 +67,18 @@ class Demo:
         self.device = 'cuda'
         self.weight_dtype = torch.bfloat16
         model_id = "stabilityai/sdxl-turbo"
+        self.current_model = 'SDXL Turbo'
         euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
         self.pipe = StableDiffusionXLPipeline.from_pretrained(model_id, scheduler=euler_anc, torch_dtype=self.weight_dtype).to(self.device)
         self.pipe.enable_xformers_memory_efficient_attention()
+        
+        self.guidance_scale = 1
+        self.num_inference_steps = 3
+        
         with gr.Blocks() as demo:
             self.layout()
             demo.queue(max_size=5).launch(share=True, max_threads=2)
-
+        
 
     def layout(self):
 
@@ -103,6 +108,13 @@ class Demo:
                         )
 
                         with gr.Row():
+                            
+                            self.model_type = gr.Dropdown(
+                                label="Model",
+                                choices= ['SDXL Turbo', 'SDXL'],
+                                value='SDXL Turbo',
+                                interactive=True
+                            )
 
                             self.model_dropdown = gr.Dropdown(
                                 label="Pretrained Sliders",
@@ -141,13 +153,14 @@ class Demo:
 
                         with gr.Row():
 
-                            self.image_new = gr.Image(
-                                label="Slider",
+                            self.image_orig = gr.Image(
+                                label="Original SD",
                                 interactive=False,
                                 type='pil',
                             )
-                            self.image_orig = gr.Image(
-                                label="Original SD",
+                            
+                            self.image_new = gr.Image(
+                                label=f"{self.model_dropdown} Slider",
                                 interactive=False,
                                 type='pil',
                             )
@@ -234,7 +247,8 @@ class Demo:
             self.seed_infr,
             self.start_noise_infr,
             self.slider_scale_infr,
-            self.model_dropdown
+            self.model_dropdown,
+            self.model
             ],
             outputs=[
                 self.image_new,
@@ -286,10 +300,25 @@ class Demo:
         return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom slider in the "Test" tab'), f'models/{save_name}', gr.update(choices=list(model_map.keys()), value=save_name.replace('.pt',''))]
 
     
-    def inference(self, prompt, seed, start_noise, scale, model_name, pbar = gr.Progress(track_tqdm=True)):
+    def inference(self, prompt, seed, start_noise, scale, model_name, model, pbar = gr.Progress(track_tqdm=True)):
         
         seed = seed or 42
-
+        if self.current_model != model:
+            if model=='SDXL Turbo':
+                model_id = "stabilityai/sdxl-turbo"
+                euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+                self.pipe = StableDiffusionXLPipeline.from_pretrained(model_id, scheduler=euler_anc, torch_dtype=self.weight_dtype).to(self.device)
+                self.pipe.enable_xformers_memory_efficient_attention()
+                self.guidance_scale = 1
+                self.num_inference_steps = 3
+                self.current_model = 'SDXL Turbo'
+            else:
+                model_id = 'stabilityai/stable-diffusion-xl-base-1.0'
+                self.pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=self.weight_dtype).to(self.device)
+                self.pipe.enable_xformers_memory_efficient_attention()
+                self.guidance_scale = 7.5
+                self.num_inference_steps = 20
+                self.current_model = 'SDXL'
         generator = torch.manual_seed(seed)
 
         model_path = model_map[model_name]
@@ -327,10 +356,10 @@ class Demo:
 
 
         generator = torch.manual_seed(seed)
-        edited_image = self.pipe(prompt, num_images_per_prompt=1, num_inference_steps=3, generator=generator, network=network, start_noise=int(start_noise), scale=float(scale), unet=unet, guidance_scale=1).images[0]
+        edited_image = self.pipe(prompt, num_images_per_prompt=1, num_inference_steps=self.num_inference_steps, generator=generator, network=network, start_noise=int(start_noise), scale=float(scale), unet=unet, guidance_scale=self.guidance_scale).images[0]
         
         generator = torch.manual_seed(seed)
-        original_image = self.pipe(prompt, num_images_per_prompt=1, num_inference_steps=3, generator=generator, network=network, start_noise=start_noise, scale=0, unet=unet, guidance_scale=1).images[0]
+        original_image = self.pipe(prompt, num_images_per_prompt=1, num_inference_steps=self.num_inference_steps, generator=generator, network=network, start_noise=start_noise, scale=0, unet=unet, guidance_scale=self.guidance_scale).images[0]
         
         del unet, network
         unet = None