crfm-helm 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

helm/clients/image_generation/cogview2/sr_pipeline/itersr_model.py

@@ -0,0 +1,269 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   itersr_model.py
+@Time    :   2021/10/02 01:36:32
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+# here put the import lib
+import math
+import torch
+import torch.nn.functional as F
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    from deepspeed.runtime.activation_checkpointing.checkpointing import get_cuda_rng_tracker
+    from SwissArmyTransformer.model.base_model import BaseModel, BaseMixin
+    from SwissArmyTransformer.mpu.utils import sqrt
+    from SwissArmyTransformer.model.transformer import split_tensor_along_last_dim
+    from SwissArmyTransformer.ops.local_attention_function import f_similar, f_weighting
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+class PositionEmbeddingMixin(BaseMixin):
+    def __init__(
+        self, additional_sequence_length, hidden_size, init_method_std=0.02, reinit_slice=slice(512, 512 + 400)
+    ):
+        super(PositionEmbeddingMixin, self).__init__()
+        self.reinit_slice = reinit_slice
+        self.position_embeddings = torch.nn.Embedding(additional_sequence_length, hidden_size)
+        torch.nn.init.normal_(self.position_embeddings.weight, mean=0.0, std=init_method_std)
+
+    def reinit(self, parent_model=None):
+        old_weights = self.transformer.position_embeddings.weight.data[self.reinit_slice]
+        old_len, hidden_size = old_weights.shape
+        assert hidden_size == self.position_embeddings.weight.shape[-1]
+        old_edge, new_edge = sqrt(old_len), sqrt(self.position_embeddings.weight.shape[-2])
+        assert new_edge % old_edge == 0
+        self.position_embeddings.weight.data.view(
+            new_edge // old_edge, old_edge, new_edge // old_edge, old_edge, hidden_size
+        ).copy_(old_weights.view(1, old_edge, 1, old_edge, hidden_size))
+
+
+class ItersrModel(BaseModel):
+    def __init__(self, args, transformer=None):
+        super().__init__(args, transformer=transformer)
+        self.original_sequence_length = args.max_sequence_length
+        additional_seqlen = args.new_sequence_length - args.max_sequence_length
+        self.add_mixin("extra_position_embedding", PositionEmbeddingMixin(additional_seqlen, args.hidden_size))
+        # self.add_mixin('attention_plus', AttentionMixin(
+        #     num_layers=args.num_layers,
+        #     hidden_size=args.hidden_size
+        # ))
+        self.layout = args.layout
+        # [PAD]... [ROI1] text ... [BOI1] {layout[0]} 1024 {layout[1]} [EOI1] 4095 {layout[2]}
+        self.kernel_size = args.kernel_size
+        self.kernel_size2 = args.kernel_size2
+        self.log_attention_weights = None
+
+    def position_embedding_forward(self, position_ids, **kw_args):
+        position = position_ids[..., : self.layout[0]]
+        position_plus = position_ids[..., self.layout[0] :] - self.original_sequence_length
+        position_embeddings = torch.cat(
+            (
+                self.transformer.position_embeddings(position),
+                self.get_mixin("extra_position_embedding").position_embeddings(position_plus),
+            ),
+            dim=-2,
+        )
+        return position_embeddings
+
+    def attention_forward(self, hidden_states, mask, layer_id=None, log_attention_weights=None, **kw_args):
+        attn_module = self.transformer.layers[layer_id].attention
+        # base model qkv
+        mixed_raw_layer = attn_module.query_key_value(hidden_states)
+        q0, k0, v0 = split_tensor_along_last_dim(mixed_raw_layer[:, : self.layout[0]], 3)
+        # cuda2d model qkv
+        q1, k1, v1 = split_tensor_along_last_dim(mixed_raw_layer[:, self.layout[0] :], 3)
+
+        dropout_fn = attn_module.attention_dropout if self.training else None
+
+        # cuda2d attention
+        context_layer = sparse_attention_2d_text(
+            q0,
+            k0,
+            v0,
+            q1,
+            k1,
+            v1,
+            mask,
+            n_head=attn_module.num_attention_heads_per_partition,
+            text_len=self.layout[0],
+            kernel_size=self.kernel_size,
+            attention_dropout=dropout_fn,
+            log_attention_weights=log_attention_weights,
+        )
+
+        output = attn_module.dense(context_layer)
+
+        return output
+
+    def final_forward(self, logits, **kwargs):
+        logits_parallel = logits
+        logits_parallel = torch.nn.functional.linear(
+            logits_parallel, self.transformer.word_embeddings.weight[:20000]
+        ).float()
+        # logits_parallel = torch.nn.functional.linear(logits_parallel, self.transformer.word_embeddings.weight[:20000])
+        return logits_parallel
+
+    # def disable_untrainable_params(self):
+    #     self.transformer.requires_grad_(False)
+
+    @classmethod
+    def add_model_specific_args(cls, parser):
+        group = parser.add_argument_group("Cuda2dModel", "cuda2d model configurations")
+        group.add_argument("--kernel-size", type=int, default=5)
+        group.add_argument("--kernel-size2", type=int, default=5)
+        group.add_argument("--layout", type=str, default="16,3616")
+        group.add_argument("--new-sequence-length", type=int, default=4096)
+        return parser
+
+
+def sparse_attention_2d_text(
+    q0,
+    k0,
+    v0,
+    q1,
+    k1,
+    v1,
+    attention_mask,
+    n_head,
+    text_len,
+    kernel_size=9,
+    attention_dropout=None,
+    log_attention_weights=None,
+    **kwargs,
+):
+    """
+    q0, k0, v0: [batch_size, 16, hidden_size]
+    q1, k1, v1: [batch_size, 3600, hidden_size]
+    n_head: int
+    attention_mask: [batch_size, 16]
+    """
+    b, s0, h0 = q0.shape
+    b, s1, h1 = q1.shape
+    h, l1 = h0 // n_head, sqrt(s1)
+    assert attention_mask.shape[-1] == s0, f"Mask Shape: {attention_mask.shape}"
+
+    q0 = q0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
+    v0 = v0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
+    k0T = k0.reshape(b, s0, n_head, h).permute(0, 2, 3, 1)
+
+    # standard attention for level 0
+    attention_scores = torch.matmul(q0 / math.sqrt(q0.shape[-1]), k0T)
+
+    attention_scores = torch.mul(attention_scores, attention_mask) - 10000.0 * (1.0 - attention_mask)
+
+    attention_probs0 = F.softmax(attention_scores, dim=-1)
+
+    # local attention for level 1
+    q1 = (
+        (q1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1) / math.sqrt(h1 // n_head))
+        .contiguous()
+        .view(b * n_head, h1 // n_head, l1, l1)
+    )
+    k1 = k1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
+    v1 = v1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
+    scores_1_to_1 = f_similar(q1, k1, kernel_size * 2 - 1, kernel_size, False)
+
+    # cross attention
+    scores_1_to_0 = torch.matmul(q1.view(b, n_head, h, s1).transpose(-1, -2), k0T)
+    if log_attention_weights is not None:
+        scores_1_to_0 += log_attention_weights
+    scores_1_to_0 = torch.mul(scores_1_to_0, attention_mask) - 10000.0 * (1.0 - attention_mask)
+    scores_1 = torch.cat(
+        (scores_1_to_0.view(b * n_head, s1, s0), scores_1_to_1.view(b * n_head, -1, scores_1_to_1.shape[3])), dim=-1
+    )
+    attention_probs1 = F.softmax(scores_1, dim=-1)
+
+    if attention_dropout is not None:
+        with get_cuda_rng_tracker().fork():
+            attention_probs1 = attention_dropout(attention_probs1)
+
+    # weighting for level 0
+    context0 = torch.matmul(attention_probs0, v0)  # [b, n_head, s0, h]
+    # weighting for level 1
+    probs_1_to_1 = attention_probs1[:, :, -scores_1_to_1.shape[3] :].view_as(scores_1_to_1)
+    context1_to_1 = f_weighting(v1, probs_1_to_1.contiguous(), kernel_size * 2 - 1, kernel_size, False)
+
+    context1 = context1_to_1.view(b, n_head, h, l1**2)
+    # weighting for cross attention
+    probs_1_to_0 = attention_probs1[:, :, : scores_1_to_0.shape[3]].view(b, n_head, -1, scores_1_to_0.shape[3])
+
+    context1_to_0 = torch.matmul(probs_1_to_0, v0)
+    context1 = context1.transpose(-1, -2) + context1_to_0
+
+    output = torch.cat((context0, context1), dim=2).transpose(1, 2).reshape(b, s0 + s1, h0)
+
+    return output
+
+
+def sparse_attention_2d_notext(
+    q0,
+    k0,
+    v0,
+    q1,
+    k1,
+    v1,
+    attention_mask,
+    n_head,
+    text_len,
+    kernel_size=9,
+    attention_dropout=None,
+    log_attention_weights=None,
+    **kwargs,
+):
+    """
+    q0, k0, v0: [batch_size, 16, hidden_size]
+    q1, k1, v1: [batch_size, 3600, hidden_size]
+    n_head: int
+    attention_mask: [batch_size, 16]
+    """
+    b, s0, h0 = q0.shape
+    b, s1, h1 = q1.shape
+    h, l1 = h0 // n_head, sqrt(s1)
+    assert len(attention_mask.shape) == 4 and attention_mask.shape[-1] == s0, f"Mask Shape: {attention_mask.shape}"
+
+    q0 = q0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
+    v0 = v0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
+    k0T = k0.reshape(b, s0, n_head, h).permute(0, 2, 3, 1)
+
+    # standard attention for level 0
+    attention_scores = torch.matmul(q0 / math.sqrt(q0.shape[-1]), k0T)
+
+    attention_scores = torch.mul(attention_scores, attention_mask) - 10000.0 * (1.0 - attention_mask)
+
+    attention_probs0 = F.softmax(attention_scores, dim=-1)
+
+    # local attention for level 1
+    q1 = (
+        (q1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1) / math.sqrt(h1 // n_head))
+        .contiguous()
+        .view(b * n_head, h1 // n_head, l1, l1)
+    )
+    k1 = k1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
+    v1 = v1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
+    scores_1_to_1 = f_similar(q1, k1, kernel_size * 2 - 1, kernel_size, False)
+
+    attention_probs1 = F.softmax(scores_1_to_1, dim=-1)
+
+    if attention_dropout is not None:
+        with get_cuda_rng_tracker().fork():
+            attention_probs1 = attention_dropout(attention_probs1)
+
+    # weighting for level 0
+    context0 = torch.matmul(attention_probs0, v0)  # [b, n_head, s0, h]
+    # weighting for level 1
+    probs_1_to_1 = attention_probs1
+    context1_to_1 = f_weighting(v1, probs_1_to_1.contiguous(), kernel_size * 2 - 1, kernel_size, False)
+
+    context1 = context1_to_1.view(b, n_head, h, l1**2)
+    # weighting for cross attention
+    context1 = context1.transpose(-1, -2)
+
+    output = torch.cat((context0, context1), dim=2).transpose(1, 2).reshape(b, s0 + s1, h0)
+
+    return output

helm/clients/image_generation/cogview2/sr_pipeline/itersr_sampling.py

@@ -0,0 +1,120 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   itersr_sampling.py
+@Time    :   2022/03/03 14:24:28
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+# here put the import lib
+import torch
+import torch.nn.functional as F
+from icetk import icetk as tokenizer
+
+
+def top_k_logits_(logits, top_k=0, filter_value=-float("Inf")):
+    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+    logits[indices_to_remove] = filter_value
+    return logits
+
+
+class IterativeEntfilterStrategy:
+    def __init__(self, invalid_slices=[], temperature=1.0, topk=10):
+        self.invalid_slices = invalid_slices
+        self.temperature = temperature
+        self.topk = topk
+
+    def forward(self, logits, tokens, temperature=None, entfilter=None, filter_topk=5, temperature2=None):
+        # In interative strategy, logits are of shape [batch_size, seq_length, hidden_size]
+        if temperature is None:
+            temperature = self.temperature
+
+        logits = logits.float() / temperature
+        for invalid_slice in self.invalid_slices:
+            logits[..., invalid_slice] = -float("Inf")
+
+        # debiased topk
+        # probs = F.softmax(logits, dim=-1)
+        # tk_value, tk_idx = torch.topk(probs, self.topk, dim=-1)
+        # pred = torch.multinomial(probs.view(-1, logits.shape[-1]), num_samples=1).view(*logits.shape[:2], 1)
+        # edge_idx = tk_idx[:, :, -1:]
+        # edge_value = tk_value[:, :, -1:]
+        # edge_mask = probs.gather(dim=-1, index=pred) < edge_value
+        # pred[edge_mask] = edge_idx[edge_mask] # replace outliers as the "filter_topk"-th token
+        # pred.squeeze_(-1) # [batch_size, seq_length]
+
+        top_k_logits_(logits, self.topk)
+        probs = F.softmax(logits, dim=-1)
+        pred = torch.multinomial(probs.view(-1, logits.shape[-1]), num_samples=1).view(*logits.shape[:2], 1)
+        pred.squeeze_(-1)
+
+        assert tokens.shape[1] == pred.shape[1]
+        tokens = pred
+        return tokens
+
+
+def filling_sequence_itersr(
+    model,
+    seq0,
+    seq1,
+    warmup_steps=3,
+    block_hw=(4, 4),
+    strategy=IterativeEntfilterStrategy(topk=10),
+):
+    """
+    seq: [PAD]... [ROI1] text ... [BOI1] {layout[0]} 1024 {layout[1]} [EOI1]
+        4095 {layout[2]} final_token.
+    Attention:
+    The sampling temperature are changing, temporally we hard code them here.
+    The temperature in the strategy is not used.
+    """
+    assert hasattr(model, "layout")
+    layout = model.layout
+
+    device = seq0.device
+    # concat and pad sequences
+    batch_size = seq0.shape[0]
+    n_pad = layout[0] - seq0.shape[1]
+    assert n_pad >= 0, "You should truncate long input before filling."
+    seq = torch.cat(
+        (torch.tensor([0] * n_pad, device=device, dtype=seq0.dtype).unsqueeze(0).expand(batch_size, n_pad), seq0, seq1),
+        dim=1,
+    )  # [b, layout[-1]+1]
+    assert seq.shape[1] == layout[-1]
+
+    # build initial tokens, attention_mask, and position_ids
+    tokens = seq.clone()
+    attention_mask = torch.ones(layout[0]).to(device)
+    attention_mask[:n_pad] = 0
+    attention_mask = attention_mask.unsqueeze(0).type_as(next(model.parameters()))  # if fp16
+    position_ids = torch.cat(
+        (
+            torch.zeros(n_pad, dtype=torch.long),
+            torch.arange(0, layout[0] - n_pad),
+            torch.arange(1024, 1024 + layout[1] - layout[0]),
+        )
+    ).to(device)
+    log_attention_weights = torch.zeros(layout[0], device=device).type_as(next(model.parameters()))
+    log_attention_weights[n_pad : layout[0]] = 0.0
+    log_attention_weights = log_attention_weights.unsqueeze(0)
+
+    # prepare for interation
+    unfixed = tokens == tokenizer["<start_of_image>"]
+    ll, rr = block_hw
+    # edge_len = int(math.sqrt(layout[-1] - layout[-2]) + 1e-4)
+    num_steps = 1
+    # interative refining
+
+    # unfixed[..., -(layout[-1] - layout[-2]):].view(
+    #     batch_size, edge_len//ll, ll, edge_len//rr, rr)[:, :, :, :, -1] = False
+
+    ret = []
+    # ret.append(tokens[:, layout[-2]:-1].clone())
+    for step_cnt in range(1, num_steps + 1):
+        logits, *_dump = model(tokens, position_ids, attention_mask, log_attention_weights=log_attention_weights)
+        real_temp = 1.0
+        new_tokens = strategy.forward(logits, tokens, real_temp)
+        tokens[unfixed] = new_tokens[unfixed]
+
+        ret.append(tokens[:, layout[-2] :].clone())
+    return torch.cat(ret, dim=0)

helm/clients/image_generation/cogview2/sr_pipeline/sr_group.py

@@ -0,0 +1,42 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   sr_group.py
+@Time    :   2022/04/02 01:17:21
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+# here put the import lib
+from .direct_sr import DirectSuperResolution
+from .iterative_sr import IterativeSuperResolution
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+
+class SRGroup:
+    def __init__(
+        self,
+        args,
+        home_path=None,
+    ):
+        try:
+            from SwissArmyTransformer.resources import auto_create
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        dsr_path = auto_create("cogview2-dsr", path=home_path)
+        itersr_path = auto_create("cogview2-itersr", path=home_path)
+        dsr = DirectSuperResolution(args, dsr_path)
+        itersr = IterativeSuperResolution(args, itersr_path, shared_transformer=dsr.model.transformer)
+        self.dsr = dsr
+        self.itersr = itersr
+
+    def sr_base(self, img_tokens, txt_tokens):
+        assert img_tokens.shape[-1] == 400 and len(img_tokens.shape) == 2
+        batch_size = img_tokens.shape[0]
+        txt_len = txt_tokens.shape[-1]
+        if len(txt_tokens.shape) == 1:
+            txt_tokens = txt_tokens.unsqueeze(0).expand(batch_size, txt_len)
+        sred_tokens = self.dsr(txt_tokens, img_tokens)
+        iter_tokens = self.itersr(txt_tokens, sred_tokens[:, -3600:].clone())
+        return iter_tokens[-batch_size:]

helm/clients/image_generation/cogview2_client.py

@@ -0,0 +1,191 @@
+import os
+import argparse
+from functools import partial
+from typing import Any, Dict, List, Optional
+
+import torch
+from icetk import icetk as tokenizer
+from torchvision.utils import save_image
+
+from helm.common.cache import CacheConfig, Cache
+from helm.common.file_caches.file_cache import FileCache
+from helm.common.hierarchical_logger import hlog, htrack_block
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.request import Request, RequestResult, GeneratedOutput, wrap_request_time
+from helm.common.tokenization_request import (
+    DecodeRequest,
+    DecodeRequestResult,
+    TokenizationRequest,
+    TokenizationRequestResult,
+)
+from helm.clients.client import Client, CachingClient
+from helm.clients.image_generation.cogview2.coglm_strategy import CoglmStrategy
+from .image_generation_client_utils import get_single_image_multimedia_object
+
+
+class CogView2Client(Client):
+    """
+    https://github.com/THUDM/CogView2
+    """
+
+    MAX_SEQ_LEN: int = 95
+    MODEL_URL: str = "https://nlp.stanford.edu/projects/vhelm/cogview2/sharefs.zip"
+
+    def __init__(self, cache_config: CacheConfig, file_cache: FileCache):
+        self._cache = Cache(cache_config)
+        self._file_cache: FileCache = file_cache
+
+        self._args: Optional[argparse.Namespace] = None
+        self._strategy: Optional[CoglmStrategy] = None
+        self._model = None
+        self._srg = None
+
+    def _get_model(self) -> None:
+        try:
+            from SwissArmyTransformer import get_args
+            from helm.clients.image_generation.cogview2.coglm_utils import (
+                get_recipe,
+                InferenceModel,
+            )
+            from helm.clients.image_generation.cogview2.sr_pipeline import SRGroup
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        tokenizer.add_special_tokens(["<start_of_image>", "<start_of_english>", "<start_of_chinese>"])
+
+        model_local_path: str = f"{self._file_cache._location}/cogview2"  # type: ignore
+        os.environ["SAT_HOME"] = f"{model_local_path}/sharefs/cogview-new"
+
+        # Download the model if not yet
+        if not os.path.exists(model_local_path):
+            os.system(f"mkdir -p {model_local_path}")
+            os.system(f"wget {self.MODEL_URL} -P {model_local_path}")
+            os.system(f"unzip {model_local_path}/sharefs.zip -d {model_local_path}")
+
+        if self._model is None:
+            # Set up args
+            args = get_args("--mode inference --fp16".split())
+            self._args = argparse.Namespace(**vars(args), **get_recipe("none"))
+            self._args.img_size = 160
+            self._args.only_first_stage = False
+            self._args.inverse_prompt = False
+            self._args.batch_size = 1
+            self._args.max_inference_batch_size = 1
+
+            # Load the model components
+            self._model, self._args = InferenceModel.from_pretrained(self._args, "coglm")
+            invalid_slices = [slice(tokenizer.num_image_tokens, None)]
+            self._strategy = CoglmStrategy(
+                invalid_slices,
+                temperature=getattr(self._args, "temp_all_gen"),
+                top_k=getattr(self._args, "topk_gen"),
+                top_k_cluster=getattr(self._args, "temp_cluster_gen"),
+            )
+            self._srg = SRGroup(self._args)  # type: ignore
+
+    def _model_inference(self, prompt) -> torch.Tensor:
+        try:
+            from SwissArmyTransformer.generation.autoregressive_sampling import filling_sequence
+            from helm.clients.image_generation.cogview2.coglm_utils import get_masks_and_position_ids_coglm
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        with torch.no_grad():
+            text = getattr(self._args, "query_template").format(prompt)
+            seq = tokenizer.encode(text)
+            if len(seq) > self.MAX_SEQ_LEN:
+                seq = seq[: self.MAX_SEQ_LEN - 2] + seq[-2:]
+            txt_len = len(seq) - 1
+            device = getattr(self._args, "device")
+            seq = torch.tensor(seq + [-1] * 400, device=device)
+            # calibrate text length
+            log_attention_weights = torch.zeros(
+                len(seq), len(seq), device=device, dtype=torch.half if getattr(self._args, "fp16") else torch.float32
+            )
+            log_attention_weights[:, :txt_len] = getattr(self._args, "attn_plus")
+            # generation
+            mbz = getattr(self._args, "max_inference_batch_size")
+            batch_size = getattr(self._args, "batch_size")
+            assert batch_size < mbz or batch_size % mbz == 0
+            get_func = partial(get_masks_and_position_ids_coglm, context_length=txt_len)
+            output_list = []
+            for tim in range(max(batch_size // mbz, 1)):
+                setattr(self._strategy, "start_pos", txt_len + 1)
+                coarse_samples = filling_sequence(
+                    self._model,
+                    seq.clone(),
+                    batch_size=min(batch_size, mbz),
+                    strategy=self._strategy,
+                    log_attention_weights=log_attention_weights,
+                    get_masks_and_position_ids=get_func,
+                )[0]
+                output_list.append(coarse_samples)
+
+            output_tokens = torch.cat(output_list, dim=0)
+            images = []
+            iter_tokens = getattr(self._srg, "sr_base")(output_tokens[:, -400:], seq[:txt_len])
+            for seq in iter_tokens:
+                decoded_img = tokenizer.decode(image_ids=seq[-3600:])
+                decoded_img = torch.nn.functional.interpolate(decoded_img, size=(480, 480))
+                images.append(decoded_img)  # only the last image (target)
+            return images[0]
+
+    def make_request(self, request: Request) -> RequestResult:
+        raw_request = {
+            "prompt": request.prompt,
+        }
+
+        try:
+
+            def do_it() -> Dict[str, Any]:
+                prompt: str = request.prompt
+
+                with htrack_block(f"Generating images for prompt: {prompt}"):
+                    self._get_model()
+
+                    images: List[torch.Tensor] = []
+                    for _ in range(request.num_completions):
+                        output = self._model_inference(**raw_request).cpu()  # (1, 3, 480, 480)
+                        images.append(output)
+
+                    assert (
+                        len(images) == request.num_completions
+                    ), f"Expected {request.num_completions} images, but got {len(images)}"
+
+                    result: Dict = {"file_locations": []}
+                    for image in images:
+                        # Write out the image to a file and save the path
+                        file_location: str = self._file_cache.generate_unique_new_file_path()  # type: ignore
+                        save_image(image, file_location, normalize=True)
+                        hlog(f"Image saved at {file_location}.")
+                        result["file_locations"].append(file_location)
+                    return result
+
+            # Include the model name and number of completions in the cache key
+            cache_key = CachingClient.make_cache_key(
+                {"model": request.model_engine, "n": request.num_completions, **raw_request}, request
+            )
+            results, cached = self._cache.get(cache_key, wrap_request_time(do_it))
+        except RuntimeError as e:
+            error: str = f"CogView2Client error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = [
+            GeneratedOutput(
+                text="", logprob=0, tokens=[], multimodal_content=get_single_image_multimedia_object(location)
+            )
+            for location in results["file_locations"]
+        ]
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=results["request_time"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
+        raise NotImplementedError("This client does not support tokenizing.")
+
+    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
+        raise NotImplementedError("This client does not support decoding.")