def predict(
        self,
        prompt: str = Input(
            description="Input Prompt. You can let language model automatically identify keywords, or provide them below.",
            default="A beautiful city skyline stamp of Shanghai",
        ),
        keywords: str = Input(
            description="(Optional) Keywords. Should be seperated by / (e.g., keyword1/keyword2/...).",
            default=None,
        ),
        positive_prompt: str = Input(
            description="(Optional) Positive prompt.",
            default=", digital art, very detailed, fantasy, high definition, cinematic light, dnd, trending on artstation",
        ),
        use_lcm: bool = Input(
            description="Use Latent Consistent Model.", default=False
        ),
        generate_natural_image: bool = Input(
            description="If set to True, the text position and content info will not be incorporated.",
            default=False,
        ),
        num_images: int = Input(
            description="Number of Output images.", default=1, ge=1, le=4
        ),
        num_inference_steps: int = Input(
            description="Number of denoising steps. You may decease the step to 4 when using LCM.",
            ge=1,
            le=50,
            default=20,
        ),
        guidance_scale: float = Input(
            description="Scale for classifier-free guidance. The scale is set to 7.5 by default. When using LCM, guidance_scale is set to 1.",
            ge=1,
            le=20,
            default=7.5,
        ),
        temperature: float = Input(
            description="Control the diversity of layout planner. Higher value indicates more diversity.",
            ge=0.1,
            le=2,
            default=1.4,
        ),
    ) -> ModelOutput:
        """Run a single prediction on the model"""
        if positive_prompt is not None and not len(positive_prompt.strip()) == 0:
            prompt += positive_prompt

        with torch.no_grad():
            user_prompt = prompt

            if generate_natural_image:
                composed_prompt = user_prompt
                prompt = self.tokenizer.encode(user_prompt)
            else:
                if keywords is None or len(keywords.strip()) == 0:
                    template = f"Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}"
                else:
                    keywords = keywords.split("/")
                    keywords = [i.strip() for i in keywords]
                    template = f"Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. In addition, we also provide all keywords at random order for reference. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {prompt}. Keywords: {str(keywords)}"

                msg = template
                conv = get_conversation_template(self.m1_model_path)
                conv.append_message(conv.roles[0], msg)
                conv.append_message(conv.roles[1], None)
                prompt = conv.get_prompt()
                inputs = self.m1_tokenizer([prompt], return_token_type_ids=False)
                inputs = {k: torch.tensor(v).to("cuda") for k, v in inputs.items()}
                output_ids = self.m1_model.generate(
                    **inputs,
                    do_sample=True,
                    temperature=temperature,
                    repetition_penalty=1.0,
                    max_new_tokens=512,
                )

                if self.m1_model.config.is_encoder_decoder:
                    output_ids = output_ids[0]
                else:
                    output_ids = output_ids[0][len(inputs["input_ids"][0]) :]

                outputs = self.m1_tokenizer.decode(
                    output_ids,
                    skip_special_tokens=True,
                    spaces_between_special_tokens=False,
                )
                print(f"[{conv.roles[0]}]\n{msg}")
                print(f"[{conv.roles[1]}]\n{outputs}")

                ocrs = outputs.split("\n")
                current_ocr = ocrs

                ocr_ids = []
                print("user_prompt", user_prompt)
                print("current_ocr", current_ocr)

                for ocr in current_ocr:
                    ocr = ocr.strip()

                    if len(ocr) == 0 or "###" in ocr or ".com" in ocr:
                        continue

                    items = ocr.split()
                    pred = " ".join(items[:-1])
                    box = items[-1]

                    l, t, r, b = box.split(",")
                    l, t, r, b = int(l), int(t), int(r), int(b)
                    ocr_ids.extend(
                        ["l" + str(l), "t" + str(t), "r" + str(r), "b" + str(b)]
                    )

                    char_list = list(pred)
                    char_list = [f"[{i}]" for i in char_list]
                    ocr_ids.extend(char_list)
                    ocr_ids.append(self.tokenizer.eos_token_id)

                caption_ids = (
                    self.tokenizer(user_prompt, truncation=True, return_tensors="pt")
                    .input_ids[0]
                    .tolist()
                )

                try:
                    ocr_ids = self.tokenizer.encode(ocr_ids)
                    prompt = caption_ids + ocr_ids
                except:
                    prompt = caption_ids

                user_prompt = self.tokenizer.decode(prompt)
                composed_prompt = self.tokenizer.decode(prompt)

            prompt = prompt[:77]
            while len(prompt) < 77:
                prompt.append(self.tokenizer.pad_token_id)

            if not use_lcm:
                prompts_cond = prompt
                prompts_nocond = [self.tokenizer.pad_token_id] * 77

                prompts_cond = [prompts_cond] * num_images
                prompts_nocond = [prompts_nocond] * num_images

                prompts_cond = torch.Tensor(prompts_cond).long().cuda()
                prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()

                scheduler = self.scheduler
                scheduler.set_timesteps(num_inference_steps)
                noise = torch.randn((num_images, 4, 64, 64)).to("cuda").half()
                input = noise

                encoder_hidden_states_cond = self.text_encoder(prompts_cond)[0].half()
                encoder_hidden_states_nocond = self.text_encoder(prompts_nocond)[
                    0
                ].half()

                for t in tqdm(scheduler.timesteps):
                    with torch.no_grad():  # classifier free guidance
                        noise_pred_cond = self.unet(
                            sample=input,
                            timestep=t,
                            encoder_hidden_states=encoder_hidden_states_cond[
                                :num_images
                            ],
                        ).sample  # b, 4, 64, 64
                        noise_pred_uncond = self.unet(
                            sample=input,
                            timestep=t,
                            encoder_hidden_states=encoder_hidden_states_nocond[
                                :num_images
                            ],
                        ).sample  # b, 4, 64, 64
                        noisy_residual = noise_pred_uncond + guidance_scale * (
                            noise_pred_cond - noise_pred_uncond
                        )  # b, 4, 64, 64
                        input = scheduler.step(noisy_residual, t, input).prev_sample
                        del noise_pred_cond
                        del noise_pred_uncond

                        torch.cuda.empty_cache()

                # decode
                input = 1 / self.vae.config.scaling_factor * input
                images = self.vae.decode(input, return_dict=False)[0]
                width, height = 512, 512
                results = []
                new_image = Image.new("RGB", (2 * width, 2 * height))
                for index, image in enumerate(images.cpu().float()):
                    image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
                    image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
                    image = Image.fromarray(
                        (image * 255).round().astype("uint8")
                    ).convert("RGB")
                    results.append(image)
                    row = index // 2
                    col = index % 2
                    new_image.paste(image, (col * width, row * height))
            else:
                generator = torch.Generator(device=self.pipe.device).manual_seed(
                    random.randint(0, 1000)
                )
                results = self.pipe(
                    prompt=user_prompt,
                    generator=generator,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=1,
                    num_images_per_prompt=num_images,
                ).images

        torch.cuda.empty_cache()
        output_paths = []

        for i, sample in enumerate(results):
            output_path = f"/tmp/out-{i}.png"
            sample.save(output_path)
            output_paths.append(Path(output_path))

        return ModelOutput(
            output_images=output_paths,
            composed_prompt=composed_prompt,
        )
Run a single prediction on the model
textdiffuser-2/predict.py