Copy files from https://github.com/hysts/CogView2_demo

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "CogView2"]
+	path = CogView2
+	url = https://github.com/THUDM/CogView2

.pre-commit-config.yaml

@@ -0,0 +1,46 @@
+exclude: ^patch
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.2.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ['--fix=lf']
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.4
+  hooks:
+  - id: docformatter
+    args: ['--in-place']
+- repo: https://github.com/pycqa/isort
+  rev: 5.10.1
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.812
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+- repo: https://github.com/google/yapf
+  rev: v0.32.0
+  hooks:
+  - id: yapf
+    args: ['--parallel', '--in-place']
+- repo: https://github.com/kynan/nbstripout
+  rev: 0.5.0
+  hooks:
+    - id: nbstripout
+      args: ['--extra-keys', 'metadata.interpreter metadata.kernelspec cell.metadata.pycharm']
+- repo: https://github.com/nbQA-dev/nbQA
+  rev: 1.3.1
+  hooks:
+    - id: nbqa-isort
+    - id: nbqa-yapf

.style.yapf

@@ -0,0 +1,5 @@
+[style]
+based_on_style = pep8
+blank_line_before_nested_class_or_def = false
+spaces_before_comment = 2
+split_before_logical_operator = true

CogView2

@@ -0,0 +1 @@
+Subproject commit 4e55cce981eb94b9c8c1f19ba9f632fd3ee42ba8

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 hysts
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

LICENSE.CogView2

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

app.py

@@ -1,32 +1,108 @@
-import os
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import argparse
+
 import gradio as gr
-os.system("pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/torch_stable.html")
-os.system("git clone https://github.com/Sleepychord/Image-Local-Attention")
-os.chdir("Image-Local-Attention")
-os.system("python setup.py install")
-os.chdir("..")
-os.system("git clone https://github.com/NVIDIA/apex")
-os.chdir("apex")
-os.system('pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./')
-os.chdir("..")
-os.system("git clone https://github.com/THUDM/CogView2")
-os.chdir("CogView2")
-os.system("gdown https://drive.google.com/uc?id=1-2nI2TTUOdiQ2WpydGafk_bZIZggQBK4")
-os.system("7za x coglm.zip")
-os.system("gdown https://drive.google.com/uc?id=1ulfXJFstYZUestvWcQIadKkNNDVbpIdM")
 
-def inference(text):
-    with open("input.txt") as f:
-        lines = f.readlines()
-    lines[0] = text
-    with open("input.txt", "w") as f:
-        f.writelines(lines)
-    os.system("python cogview2_text2image.py --mode inference --fp16 --input-source input.txt --output-path samples_sat_v0.2 --batch-size 4 --max-inference-batch-size 8 --only-first-stage")
-    return "/content/CogView2/output.png"
+from model import AppModel
+
+DESCRIPTION = '''# CogView2 (text2image)
+
+This is an unofficial demo for <a href="https://github.com/THUDM/CogView2">https://github.com/THUDM/CogView2</a>.
+
+[This Space](https://huggingface.co/spaces/chinhon/translation_eng2ch) is used for translation from English to Chinese.
+'''
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--only-first-stage', action='store_true')
+    parser.add_argument('--share', action='store_true')
+    return parser.parse_args()
+
+
+def set_example_text(example: list) -> dict:
+    return gr.Textbox.update(value=example[0])
+
+
+def main():
+    args = parse_args()
+    model = AppModel(args.only_first_stage)
+
+    with gr.Blocks(css='style.css') as demo:
+        gr.Markdown(DESCRIPTION)
 
-gr.Interface(inference,"text","image",title="CogView 2").launch(debug=True,enable_queue=True)
+        with gr.Row():
+            with gr.Column():
+                with gr.Group():
+                    text = gr.Textbox(label='Input Text')
+                    translate = gr.Checkbox(label='Translate to Chinese',
+                                            value=False)
+                    style = gr.Dropdown(choices=[
+                        'mainbody',
+                        'photo',
+                        'flat',
+                        'comics',
+                        'oil',
+                        'sketch',
+                        'isometric',
+                        'chinese',
+                        'watercolor',
+                    ],
+                                        label='Style')
+                    seed = gr.Slider(0,
+                                     100000,
+                                     step=1,
+                                     value=1234,
+                                     label='Seed')
+                    only_first_stage = gr.Checkbox(
+                        label='Only First Stage',
+                        value=args.only_first_stage,
+                        visible=not args.only_first_stage)
+                    num_images = gr.Slider(1,
+                                           16,
+                                           step=1,
+                                           value=8,
+                                           label='Number of Images')
+                    with open('samples.txt') as f:
+                        samples = [[line.strip()] for line in f.readlines()]
+                    examples = gr.Dataset(components=[text], samples=samples)
+                    run_button = gr.Button('Run')
 
+            with gr.Column():
+                with gr.Group():
+                    translated_text = gr.Textbox(label='Translated Text')
+                    with gr.Tabs():
+                        with gr.TabItem('Output (Grid View)'):
+                            result_grid = gr.Image(show_label=False)
+                        with gr.TabItem('Output (Gallery)'):
+                            result_gallery = gr.Gallery(show_label=False)
 
+        run_button.click(fn=model.run_with_translation,
+                         inputs=[
+                             text,
+                             translate,
+                             style,
+                             seed,
+                             only_first_stage,
+                             num_images,
+                         ],
+                         outputs=[
+                             translated_text,
+                             result_grid,
+                             result_gallery,
+                         ])
+        examples.click(fn=set_example_text,
+                       inputs=examples,
+                       outputs=examples.components)
 
+    demo.launch(
+        enable_queue=True,
+        share=args.share,
+    )
 
 
+if __name__ == '__main__':
+    main()

model.py

@@ -0,0 +1,377 @@
+#This code is adapted from https://github.com/THUDM/CogView2/blob/4e55cce981eb94b9c8c1f19ba9f632fd3ee42ba8/cogview2_text2image.py
+
+from __future__ import annotations
+
+import argparse
+import functools
+import logging
+import pathlib
+import sys
+import time
+from typing import Any
+
+import gradio as gr
+import numpy as np
+import torch
+from icetk import IceTokenizer
+from SwissArmyTransformer import get_args
+from SwissArmyTransformer.arguments import set_random_seed
+from SwissArmyTransformer.generation.autoregressive_sampling import \
+    filling_sequence
+from SwissArmyTransformer.model import CachedAutoregressiveModel
+
+app_dir = pathlib.Path(__file__).parent
+submodule_dir = app_dir / 'CogView2'
+sys.path.insert(0, submodule_dir.as_posix())
+
+from coglm_strategy import CoglmStrategy
+from sr_pipeline import SRGroup
+
+formatter = logging.Formatter(
+    '[%(asctime)s] %(name)s %(levelname)s: %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S')
+stream_handler = logging.StreamHandler(stream=sys.stdout)
+stream_handler.setLevel(logging.DEBUG)
+stream_handler.setFormatter(formatter)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.propagate = False
+logger.addHandler(stream_handler)
+
+ICETK_MODEL_DIR = app_dir / 'icetk_models'
+
+
+def get_masks_and_position_ids_coglm(
+        seq: torch.Tensor, context_length: int
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    tokens = seq.unsqueeze(0)
+
+    attention_mask = torch.ones((1, len(seq), len(seq)), device=tokens.device)
+    attention_mask.tril_()
+    attention_mask[..., :context_length] = 1
+    attention_mask.unsqueeze_(1)
+
+    position_ids = torch.zeros(len(seq),
+                               device=tokens.device,
+                               dtype=torch.long)
+    torch.arange(0, context_length, out=position_ids[:context_length])
+    torch.arange(512,
+                 512 + len(seq) - context_length,
+                 out=position_ids[context_length:])
+
+    position_ids = position_ids.unsqueeze(0)
+    return tokens, attention_mask, position_ids
+
+
+class InferenceModel(CachedAutoregressiveModel):
+    def final_forward(self, logits, **kwargs):
+        logits_parallel = logits
+        logits_parallel = torch.nn.functional.linear(
+            logits_parallel.float(),
+            self.transformer.word_embeddings.weight[:20000].float())
+        return logits_parallel
+
+
+def get_recipe(name: str) -> dict[str, Any]:
+    r = {
+        'attn_plus': 1.4,
+        'temp_all_gen': 1.15,
+        'topk_gen': 16,
+        'temp_cluster_gen': 1.,
+        'temp_all_dsr': 1.5,
+        'topk_dsr': 100,
+        'temp_cluster_dsr': 0.89,
+        'temp_all_itersr': 1.3,
+        'topk_itersr': 16,
+        'query_template': '{}<start_of_image>',
+    }
+    if name == 'none':
+        pass
+    elif name == 'mainbody':
+        r['query_template'] = '{} 高清摄影 隔绝<start_of_image>'
+
+    elif name == 'photo':
+        r['query_template'] = '{} 高清摄影<start_of_image>'
+
+    elif name == 'flat':
+        r['query_template'] = '{} 平面风格<start_of_image>'
+        # r['attn_plus'] = 1.8
+        # r['temp_cluster_gen'] = 0.75
+        r['temp_all_gen'] = 1.1
+        r['topk_dsr'] = 5
+        r['temp_cluster_dsr'] = 0.4
+
+        r['temp_all_itersr'] = 1
+        r['topk_itersr'] = 5
+    elif name == 'comics':
+        r['query_template'] = '{} 漫画 隔绝<start_of_image>'
+        r['topk_dsr'] = 5
+        r['temp_cluster_dsr'] = 0.4
+        r['temp_all_gen'] = 1.1
+        r['temp_all_itersr'] = 1
+        r['topk_itersr'] = 5
+    elif name == 'oil':
+        r['query_template'] = '{} 油画风格<start_of_image>'
+        pass
+    elif name == 'sketch':
+        r['query_template'] = '{} 素描风格<start_of_image>'
+        r['temp_all_gen'] = 1.1
+    elif name == 'isometric':
+        r['query_template'] = '{} 等距矢量图<start_of_image>'
+        r['temp_all_gen'] = 1.1
+    elif name == 'chinese':
+        r['query_template'] = '{} 水墨国画<start_of_image>'
+        r['temp_all_gen'] = 1.12
+    elif name == 'watercolor':
+        r['query_template'] = '{} 水彩画风格<start_of_image>'
+    return r
+
+
+def get_default_args() -> argparse.Namespace:
+    arg_list = ['--mode', 'inference', '--fp16']
+    args = get_args(arg_list)
+    known = argparse.Namespace(img_size=160,
+                               only_first_stage=False,
+                               inverse_prompt=False,
+                               style='mainbody')
+    args = argparse.Namespace(**vars(args), **vars(known),
+                              **get_recipe(known.style))
+    return args
+
+
+class Model:
+    def __init__(self, only_first_stage: bool = False):
+        self.args = get_default_args()
+        self.args.only_first_stage = only_first_stage
+
+        self.tokenizer = self.load_tokenizer()
+
+        self.model, self.args = self.load_model()
+        self.strategy = self.load_strategy()
+        self.srg = self.load_srg()
+
+        self.query_template = self.args.query_template
+        self.style = self.args.style
+        self.device = torch.device(self.args.device)
+        self.fp16 = self.args.fp16
+        self.max_batch_size = self.args.max_inference_batch_size
+        self.only_first_stage = self.args.only_first_stage
+
+    def load_tokenizer(self) -> IceTokenizer:
+        logger.info('--- load_tokenizer ---')
+        start = time.perf_counter()
+
+        tokenizer = IceTokenizer(ICETK_MODEL_DIR.as_posix())
+        tokenizer.add_special_tokens(
+            ['<start_of_image>', '<start_of_english>', '<start_of_chinese>'])
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+        return tokenizer
+
+    def load_model(self) -> tuple[InferenceModel, argparse.Namespace]:
+        logger.info('--- load_model ---')
+        start = time.perf_counter()
+
+        model, args = InferenceModel.from_pretrained(self.args, 'coglm')
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+        return model, args
+
+    def load_strategy(self) -> CoglmStrategy:
+        logger.info('--- load_strategy ---')
+        start = time.perf_counter()
+
+        invalid_slices = [slice(self.tokenizer.num_image_tokens, None)]
+        strategy = CoglmStrategy(invalid_slices,
+                                 temperature=self.args.temp_all_gen,
+                                 top_k=self.args.topk_gen,
+                                 top_k_cluster=self.args.temp_cluster_gen)
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+        return strategy
+
+    def load_srg(self) -> SRGroup:
+        logger.info('--- load_srg ---')
+        start = time.perf_counter()
+
+        srg = None if self.args.only_first_stage else SRGroup(self.args)
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+        return srg
+
+    def update_style(self, style: str) -> None:
+        if style == self.style:
+            return
+        logger.info('--- update_style ---')
+        start = time.perf_counter()
+
+        self.args = argparse.Namespace(**(vars(self.args) | get_recipe(style)))
+        self.query_template = self.args.query_template
+        logger.info(f'{self.query_template=}')
+
+        self.strategy.temperature = self.args.temp_all_gen
+
+        if self.srg is not None:
+            self.srg.dsr.strategy.temperature = self.args.temp_all_dsr
+            self.srg.dsr.strategy.topk = self.args.topk_dsr
+            self.srg.dsr.strategy.temperature2 = self.args.temp_cluster_dsr
+
+            self.srg.itersr.strategy.temperature = self.args.temp_all_itersr
+            self.srg.itersr.strategy.topk = self.args.topk_itersr
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+
+    def run(self, text: str, style: str, seed: int, only_first_stage: bool,
+            num: int) -> list[np.ndarray] | None:
+        set_random_seed(seed)
+        seq, txt_len = self.preprocess_text(text)
+        if seq is None:
+            return None
+        self.update_style(style)
+        self.only_first_stage = only_first_stage
+        tokens = self.generate_tokens(seq, txt_len, num)
+        res = self.generate_images(seq, txt_len, tokens)
+        return res
+
+    @torch.inference_mode()
+    def preprocess_text(
+            self, text: str) -> tuple[torch.Tensor, int] | tuple[None, None]:
+        logger.info('--- preprocess_text ---')
+        start = time.perf_counter()
+
+        text = self.query_template.format(text)
+        logger.info(f'{text=}')
+        seq = self.tokenizer.encode(text)
+        logger.info(f'{len(seq)=}')
+        if len(seq) > 110:
+            logger.info('The input text is too long.')
+            return None, None
+        txt_len = len(seq) - 1
+        seq = torch.tensor(seq + [-1] * 400, device=self.device)
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+        return seq, txt_len
+
+    @torch.inference_mode()
+    def generate_tokens(self,
+                        seq: torch.Tensor,
+                        txt_len: int,
+                        num: int = 8) -> torch.Tensor:
+        logger.info('--- generate_tokens ---')
+        start = time.perf_counter()
+
+        # calibrate text length
+        log_attention_weights = torch.zeros(
+            len(seq),
+            len(seq),
+            device=self.device,
+            dtype=torch.half if self.fp16 else torch.float32)
+        log_attention_weights[:, :txt_len] = self.args.attn_plus
+        get_func = functools.partial(get_masks_and_position_ids_coglm,
+                                     context_length=txt_len)
+
+        output_list = []
+        remaining = num
+        for _ in range((num + self.max_batch_size - 1) // self.max_batch_size):
+            self.strategy.start_pos = txt_len + 1
+            coarse_samples = filling_sequence(
+                self.model,
+                seq.clone(),
+                batch_size=min(remaining, self.max_batch_size),
+                strategy=self.strategy,
+                log_attention_weights=log_attention_weights,
+                get_masks_and_position_ids=get_func)[0]
+            output_list.append(coarse_samples)
+            remaining -= self.max_batch_size
+        output_tokens = torch.cat(output_list, dim=0)
+        logger.info(f'{output_tokens.shape=}')
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+        return output_tokens
+
+    @staticmethod
+    def postprocess(tensor: torch.Tensor) -> np.ndarray:
+        return tensor.cpu().mul(255).add_(0.5).clamp_(0, 255).permute(
+            1, 2, 0).to(torch.uint8).numpy()
+
+    @torch.inference_mode()
+    def generate_images(self, seq: torch.Tensor, txt_len: int,
+                        tokens: torch.Tensor) -> list[np.ndarray]:
+        logger.info('--- generate_images ---')
+        start = time.perf_counter()
+
+        logger.info(f'{self.only_first_stage=}')
+        res = []
+        if self.only_first_stage:
+            for i in range(len(tokens)):
+                seq = tokens[i]
+                decoded_img = self.tokenizer.decode(image_ids=seq[-400:])
+                decoded_img = torch.nn.functional.interpolate(decoded_img,
+                                                              size=(480, 480))
+                decoded_img = self.postprocess(decoded_img[0])
+                res.append(decoded_img)  # only the last image (target)
+        else:  # sr
+            iter_tokens = self.srg.sr_base(tokens[:, -400:], seq[:txt_len])
+            for seq in iter_tokens:
+                decoded_img = self.tokenizer.decode(image_ids=seq[-3600:])
+                decoded_img = torch.nn.functional.interpolate(decoded_img,
+                                                              size=(480, 480))
+                decoded_img = self.postprocess(decoded_img[0])
+                res.append(decoded_img)  # only the last image (target)
+
+        elapsed = time.perf_counter() - start
+        logger.info(f'Elapsed: {elapsed}')
+        logger.info('--- done ---')
+        return res
+
+
+class AppModel(Model):
+    def __init__(self, only_first_stage: bool):
+        super().__init__(only_first_stage)
+        self.translator = gr.Interface.load(
+            'spaces/chinhon/translation_eng2ch')
+
+    def make_grid(self, images: list[np.ndarray] | None) -> np.ndarray | None:
+        if images is None or len(images) == 0:
+            return None
+        ncols = 1
+        while True:
+            if ncols**2 >= len(images):
+                break
+            ncols += 1
+        nrows = (len(images) + ncols - 1) // ncols
+        h, w = images[0].shape[:2]
+        grid = np.zeros((h * nrows, w * ncols, 3), dtype=np.uint8)
+        for i in range(nrows):
+            for j in range(ncols):
+                index = ncols * i + j
+                if index >= len(images):
+                    break
+                grid[h * i:h * (i + 1), w * j:w * (j + 1)] = images[index]
+        return grid
+
+    def run_with_translation(
+        self, text: str, translate: bool, style: str, seed: int,
+        only_first_stage: bool, num: int
+    ) -> tuple[str | None, np.ndarray | None, list[np.ndarray] | None]:
+        if translate:
+            text = translated_text = self.translator(text)
+        else:
+            translated_text = None
+        results = self.run(text, style, seed, only_first_stage, num)
+        grid_image = self.make_grid(results)
+        return translated_text, grid_image, results

patch

@@ -0,0 +1,51 @@
+diff --git a/coglm_strategy.py b/coglm_strategy.py
+index cba87ce..40e4ece 100755
+--- a/coglm_strategy.py
++++ b/coglm_strategy.py
+@@ -8,6 +8,7 @@
+ 
+ # here put the import lib
+ import os
++import pathlib
+ import sys
+ import math
+ import random
+@@ -57,7 +58,8 @@ class CoglmStrategy:
+         self._is_done = False
+         self.outlier_count_down = 5
+         self.vis_list = [[]for i in range(16)]
+-        self.cluster_labels = torch.tensor(np.load('cluster_label.npy'), device='cuda', dtype=torch.long)
++        cluster_label_path = pathlib.Path(__file__).parent / 'cluster_label.npy'
++        self.cluster_labels = torch.tensor(np.load(cluster_label_path), device='cuda', dtype=torch.long)
+         self.top_k_cluster = top_k_cluster
+ 
+     @property
+@@ -91,4 +93,4 @@ class CoglmStrategy:
+ 
+     def finalize(self, tokens, mems):
+         self._is_done = False
+-        return tokens, mems
+
++        return tokens, mems
+diff --git a/sr_pipeline/dsr_sampling.py b/sr_pipeline/dsr_sampling.py
+index a0d0298..f721573 100755
+--- a/sr_pipeline/dsr_sampling.py
++++ b/sr_pipeline/dsr_sampling.py
+@@ -8,6 +8,7 @@
+ 
+ # here put the import lib
+ import os
++import pathlib
+ import sys
+ import math
+ import random
+@@ -27,7 +28,8 @@ class IterativeEntfilterStrategy:
+         self.invalid_slices = invalid_slices
+         self.temperature = temperature
+         self.topk = topk        
+-        self.cluster_labels = torch.tensor(np.load('cluster_label.npy'), device='cuda', dtype=torch.long)
++        cluster_label_path = pathlib.Path(__file__).parents[1] / 'cluster_label.npy'
++        self.cluster_labels = torch.tensor(np.load(cluster_label_path), device='cuda', dtype=torch.long)
+         self.temperature2 = temperature2
+ 
+ 

requirements.txt

@@ -1,3 +1,7 @@
-SwissArmyTransformer>=0.2
-icetk
-gdown
+git+https://github.com/Sleepychord/Image-Local-Attention@43fee31
+gradio==3.0.17
+icetk==0.0.3
+numpy==1.22.4
+SwissArmyTransformer==0.2.4
+torch==1.11.0
+torchvision==0.12.0

samples.txt

@@ -0,0 +1,13 @@
+A beautiful girl is hugging a husky.
+A lion teacher wearing a suit is in front of a blackboard.
+A robot is riding under the blue and cloudy sky.
+Several youths are talking in a bar.
+A lion man is typing in the office.
+A young woman is taking photos.
+A pirate captain with a skull.
+A girl holding an oil-paper umbrella in a rainy lane.
+Earth in the eye.
+A magnificent church. Sketch.
+Mount Fuji, cherry blossom and Akita dog. Oil painting.
+A tiger with angel's wings.
+A fox is sitting on the books.

style.css

@@ -0,0 +1,7 @@
+h1 {
+  text-align: center;
+}
+img#visitor-badge {
+  display: block;
+  margin: auto;
+}