Update app.py

app.py

@@ -129,30 +129,39 @@ def process(
         shape = (1, 4, height // 8, width // 8)
         x_T = torch.randn(shape, device=device, dtype=torch.float32)
 
+        # Modify the get_learned_conditioning method to handle the attention mask issue
+        def modified_get_learned_conditioning(model, prompt):
+            tokens = model.cond_stage_model.tokenizer.encode(prompt)
+            tokens = torch.LongTensor(tokens).to(model.device).unsqueeze(0)
+            encoder_hidden_states = model.cond_stage_model.transformer(input_ids=tokens).last_hidden_state
+            return encoder_hidden_states
+
+        cond = modified_get_learned_conditioning(model, positive_prompt)
+        uncond = modified_get_learned_conditioning(model, negative_prompt)
+
         if not tile_diffusion and not tile_vae:
             samples = sampler.sample_ccsr(
                 steps=steps, t_max=0.6667, t_min=0.3333, shape=shape, cond_img=control,
-                positive_prompt=positive_prompt, negative_prompt=negative_prompt, x_T=x_T,
+                positive_prompt=cond, negative_prompt=uncond, x_T=x_T,
                 cfg_scale=cfg_scale,
                 color_fix_type="adain" if use_color_fix else "none"
             )
         else:
             if tile_vae:
-                # Remove this line as ControlLDM doesn't have _init_tiled_vae method
-                # model._init_tiled_vae(encoder_tile_size=vae_encoder_tile_size, decoder_tile_size=vae_decoder_tile_size)
+                # Note: Tiled VAE is not implemented in this version
                 pass
             if tile_diffusion:
                 samples = sampler.sample_with_tile_ccsr(
                     tile_size=tile_diffusion_size, tile_stride=tile_diffusion_stride,
                     steps=steps, t_max=0.6667, t_min=0.3333, shape=shape, cond_img=control,
-                    positive_prompt=positive_prompt, negative_prompt=negative_prompt, x_T=x_T,
+                    positive_prompt=cond, negative_prompt=uncond, x_T=x_T,
                     cfg_scale=cfg_scale,
                     color_fix_type="adain" if use_color_fix else "none"
                 )
             else:
                 samples = sampler.sample_ccsr(
                     steps=steps, t_max=0.6667, t_min=0.3333, shape=shape, cond_img=control,
-                    positive_prompt=positive_prompt, negative_prompt=negative_prompt, x_T=x_T,
+                    positive_prompt=cond, negative_prompt=uncond, x_T=x_T,
                     cfg_scale=cfg_scale,
                     color_fix_type="adain" if use_color_fix else "none"
                 )