update

app.py

@@ -59,13 +59,13 @@ m1_model_path = 'JingyeChen22/textdiffuser2_layout_planner'
 
 m1_tokenizer = AutoTokenizer.from_pretrained(m1_model_path, use_fast=False)
 m1_model = AutoModelForCausalLM.from_pretrained(
-    m1_model_path, low_cpu_mem_usage=True
+    m1_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
 ).cuda()
 
 #### import diffusion models
 text_encoder = CLIPTextModel.from_pretrained(
     'JingyeChen22/textdiffuser2-full-ft', subfolder="text_encoder", ignore_mismatched_sizes=True
-).cuda()
+).cuda().half()
 tokenizer = CLIPTokenizer.from_pretrained(
     'runwayml/stable-diffusion-v1-5', subfolder="tokenizer"
 )
@@ -83,20 +83,20 @@ for c in alphabet:
 print(len(tokenizer))
 print('***************')
 
-vae = AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="vae").cuda()
+vae = AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="vae").half().cuda()
 unet = UNet2DConditionModel.from_pretrained(
     'JingyeChen22/textdiffuser2-full-ft', subfolder="unet"
-).cuda()
+).half().cuda()
 text_encoder.resize_token_embeddings(len(tokenizer))
 
 
-# #### load lcm components
-# model_id = "lambdalabs/sd-pokemon-diffusers"
-# lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
-# pipe = DiffusionPipeline.from_pretrained(model_id, unet=unet, tokenizer=tokenizer, text_encoder=text_encoder)
-# pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-# pipe.load_lora_weights(lcm_lora_id)
-# pipe.to(device="cuda")
+#### load lcm components
+model_id = "lambdalabs/sd-pokemon-diffusers"
+lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
+pipe = DiffusionPipeline.from_pretrained(model_id, unet=unet, tokenizer=tokenizer, text_encoder=text_encoder, torch_dtype=torch.float16)
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+pipe.load_lora_weights(lcm_lora_id)
+pipe.to(device="cuda")
 
 
 #### for interactive
@@ -339,11 +339,11 @@ def text_to_image(prompt,keywords,radio,slider_step,slider_guidance,slider_batch
 
             scheduler = DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="scheduler") 
             scheduler.set_timesteps(slider_step) 
-            noise = torch.randn((slider_batch, 4, 64, 64)).to("cuda") 
+            noise = torch.randn((slider_batch, 4, 64, 64)).to("cuda").half()
             input = noise
 
-            encoder_hidden_states_cond = text_encoder(prompts_cond)[0]
-            encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0]
+            encoder_hidden_states_cond = text_encoder(prompts_cond)[0].half()
+            encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0].half()
 
 
             for t in tqdm(scheduler.timesteps):