Update app.py

app.py

@@ -20,7 +20,7 @@ import paramiko
 import datetime
 #from diffusers import DPMSolverSDEScheduler
 from diffusers.models.attention_processor import AttnProcessor2_0
-from transformers import CLIPTextModelWithProjection, CLIPTextModel
+from transformers import CLIPTextModelWithProjection, CLIPTextModel, CLIPTokenizer
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
@@ -82,11 +82,17 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 os.putenv("HF_HUB_ENABLE_HF_TRANSFER","1")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    
+
+text_encoder=CLIPTextModel.from_pretrained('ford442/RealVisXL_V5.0_BF16', low_cpu_mem_usage=False, subfolder='text_encoder',token=True)#.to(device=device, dtype=torch.bfloat16)
+text_encoder_2=CLIPTextModelWithProjection.from_pretrained('ford442/RealVisXL_V5.0_BF16', low_cpu_mem_usage=False, subfolder='text_encoder_2',token=True)#.to(device=device, dtype=torch.bfloat16)
+tokenizer_1=CLIPTokenizer.from_pretrained('ford442/RealVisXL_V5.0_BF16', low_cpu_mem_usage=False, subfolder='tokenizer',token=True)
+tokenizer_2=CLIPTokenizer.from_pretrained('ford442/RealVisXL_V5.0_BF16', low_cpu_mem_usage=False, subfolder='tokenizer_2',token=True)
+scheduler=EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', low_cpu_mem_usage=False, subfolder='scheduler',token=True)
+vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", low_cpu_mem_usage=False, safety_checker=None, use_safetensors=False, torch_dtype=torch.float32, token=True) #.to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
+
 def load_and_prepare_model():
     #vaeRV = AutoencoderKL.from_pretrained("SG161222/RealVisXL_V5.0", subfolder='vae', safety_checker=None, use_safetensors=True, token=True)
     #vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False, low_cpu_mem_usage=False, torch_dtype=torch.float32, token=True) #.to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
-    vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", low_cpu_mem_usage=False, safety_checker=None, use_safetensors=False, torch_dtype=torch.float32, token=True) #.to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
     #sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1,use_karras_sigmas=True)
     #sched = DPMSolverSDEScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler')
     #sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear", token=True) #, beta_start=0.00085, beta_end=0.012, steps_offset=1,use_karras_sigmas=True, token=True)
@@ -96,12 +102,16 @@ def load_and_prepare_model():
         #torch_dtype=torch.bfloat16,
         token=True,
         add_watermarker=False,
-        #text_encoder=None,
-        #text_encoder_2=None,
+        text_encoder=None,
+        text_encoder_2=None,
+        tokenizer=None,
+        tokenizer_2=None,
+        scheduler=None,
         vae=None,
     )
-    #pipe.vae = vaeXL #.to(torch.bfloat16)
-    #pipe.scheduler = sched
+    pipe.scheduler=scheduler
+    pipe.tokenizer=tokenizer_1
+    pipe.tokenizer_2=tokenizer_2
     #pipe.vae.do_resize=False
     #pipe.vae.vae_scale_factor=8
     #pipe.to(device)
@@ -110,22 +120,14 @@ def load_and_prepare_model():
     pipe.watermark=None
     pipe.safety_checker=None
     #pipe.unet.to(memory_format=torch.channels_last)
-    #pipe.enable_vae_tiling()
     pipe.to(device=device, dtype=torch.bfloat16)
     pipe.vae = vaeXL.to(device) #.to('cpu') #.to(torch.bfloat16)
-    
-    # Enable attention slicing
-    pipe.enable_attention_slicing()
-
     pipe.unet.set_attn_processor(AttnProcessor2_0())
     pipe.vae.set_default_attn_processor()
     return pipe
     
 pipe = load_and_prepare_model()
 
-text_encoder=CLIPTextModel.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='text_encoder',token=True)#.to(device=device, dtype=torch.bfloat16)
-text_encoder_2=CLIPTextModelWithProjection.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='text_encoder_2',token=True)#.to(device=device, dtype=torch.bfloat16)
-   
 MAX_SEED = np.iinfo(np.int32).max
 
 neg_prompt_2 = " 'non-photorealistic':1.5, 'unrealistic skin','unattractive face':1.3, 'low quality':1.1, ('dull color scheme', 'dull colors', 'digital noise':1.2),'amateurish', 'poorly drawn face':1.3, 'poorly drawn', 'distorted face', 'low resolution', 'simplistic' "