Update app.py

app.py

@@ -12,6 +12,7 @@ from diffusers import (
     LCMScheduler, 
     AutoPipelineForText2Image,
     DPMSolverMultistepScheduler,
+    AutoencoderKL,
 )
 from transformers import pipeline
 from diffusers.utils import load_image, make_image_grid
@@ -134,6 +135,7 @@ def infer(
     use_LCM_adapter=False,   # Параметр для включения LCM_adapter
     LCM_adapter=None,        # Параметр для выбора типа LCM_adapter
     use_DDIMScheduler=False, # Параметр для включения DDIMScheduler
+    use_Tiny_VAE=False,      # Параметр для включения Tiny_VAE
     progress=gr.Progress(track_tqdm=True)
 ):
 
@@ -557,40 +559,68 @@ def infer(
 
                     image = pipe_DDIMS(**params).images[0]
                 else:    
-                    # Генерация изображений с LORA без ControlNet и IP_Adapter ---------------------------------------------------------------------------------------------
+                    # Генерация изображений с Tiny_VAE ---------------------------------------------------------------------------------------------
 
-                    print('Генерация изображений с LORA без ControlNet и IP_Adapter')
-                    
-                    # Инициализация ControlNet
-                    controlnet_model_path = "lllyasviel/sd-controlnet-openpose"
-                    controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype)
+                    if use_Tiny_VAE:
+                        
+                        print('use_Tiny_VAE = ', use_Tiny_VAE)
 
-                    generator = torch.Generator(device).manual_seed(seed)
+                        generator = torch.Generator(device).manual_seed(seed)
 
-                    if model != model_default:
-                        pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device)
-                        prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder)
-                        negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder)
-                        prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
-                    else:
-                        pipe_default = get_lora_sd_pipeline(lora_dir='lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device)
-                        pipe = pipe_default
-                        prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder)
-                        negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+                        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch_dtype)
+
+                        pipe_Tiny_VAE = StableDiffusionPipeline.from_pretrained(model_default, vae=vae, torch_dtype=torch_dtype).to(device)
+
+                        prompt_embeds = long_prompt_encoder(prompt, pipe_Tiny_VAE.tokenizer, pipe_Tiny_VAE.text_encoder)
+                        negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe_Tiny_VAE.tokenizer, pipe_Tiny_VAE.text_encoder)
                         prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
-                        pipe.fuse_lora(lora_scale=lora_scale)
 
-                    params = {
-                        'prompt_embeds': prompt_embeds,
-                        'negative_prompt_embeds': negative_prompt_embeds,
-                        'guidance_scale': guidance_scale,
-                        'num_inference_steps': num_inference_steps,
-                        'width': width,
-                        'height': height,
-                        'generator': generator,
-                    }
+                        params = {
+                            'prompt_embeds': prompt_embeds,
+                            'negative_prompt_embeds': negative_prompt_embeds,
+                            'guidance_scale': guidance_scale,
+                            'num_inference_steps': num_inference_steps,
+                            'width': width,
+                            'height': height,
+                            'generator': generator,
+                        }
+
+                        image = pipe_Tiny_VAE(**params).images[0]
+                    else:    
+                        # Генерация изображений с LORA без ControlNet и IP_Adapter ---------------------------------------------------------------------------------------------
+
+                        print('Генерация изображений с LORA без ControlNet и IP_Adapter')
+                    
+                        # Инициализация ControlNet
+                        controlnet_model_path = "lllyasviel/sd-controlnet-openpose"
+                        controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch_dtype)
+    
+                        generator = torch.Generator(device).manual_seed(seed)
 
-                    image = pipe(**params).images[0]
+                        if model != model_default:
+                            pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device)
+                            prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder)
+                            negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+                            prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+                        else:
+                            pipe_default = get_lora_sd_pipeline(lora_dir='lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device)
+                            pipe = pipe_default
+                            prompt_embeds = long_prompt_encoder(prompt, pipe.tokenizer, pipe.text_encoder)
+                            negative_prompt_embeds = long_prompt_encoder(negative_prompt, pipe.tokenizer, pipe.text_encoder)
+                            prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)
+                            pipe.fuse_lora(lora_scale=lora_scale)
+
+                        params = {
+                            'prompt_embeds': prompt_embeds,
+                            'negative_prompt_embeds': negative_prompt_embeds,
+                            'guidance_scale': guidance_scale,
+                            'num_inference_steps': num_inference_steps,
+                            'width': width,
+                            'height': height,
+                            'generator': generator,
+                        }
+
+                        image = pipe(**params).images[0]
 
     # Если выбрано удаление фона
     if remove_bg:
@@ -847,6 +877,15 @@ with gr.Blocks(css=css) as demo:
                 interactive=True
             )
 
+        # Tiny_VAE -----------------------------------------------------------------------------------------------------
+        # Checkbox для Tiny_VAE
+        with gr.Blocks():
+            use_Tiny_VAE = gr.Checkbox(
+                label="Use Tiny_VAE",
+                value=False,
+                interactive=True
+            )
+        
         # Удаление фона------------------------------------------------------------------------------------------------
         # Checkbox для удаления фона
         with gr.Blocks():
@@ -893,6 +932,7 @@ with gr.Blocks(css=css) as demo:
             use_LCM_adapter,    # Параметр для включения LCM_adapter
             LCM_adapter,        # Параметр для выбора типа LCM_adapter
             use_DDIMScheduler,  # Параметр для включения DDIMScheduler
+            use_Tiny_VAE,       # Параметр для включения Tiny_VAE
         ],
         outputs=[result],
     )