Revert "fix with main?"

This reverts commit 846e9dc, reversing
changes made to e267385.

Author: zRzRzRzRzRzRzR

.env.template

@@ -6,12 +6,3 @@ OFFLOAD_TYPE=cpu_model_offload
 COGVIEW4_PATH=THUDM/CogView4-6B
 # Optional, only needed when you don't want to use the default transformer in COGVIEW4_PATH
 # COGVIEW4_TRANSFORMER_PATH=
-
-### OPENAI API
-
-OPENAI_API_KEY=
-OPENAI_BASE_URL=
-
-### LORA PATH FOR GUI(test)
-
-LORA_DIR=

src/cogkit/api/python/generation/util.py

@@ -6,13 +6,11 @@
     CogVideoXDPMScheduler,
     CogVideoXImageToVideoPipeline,
     CogVideoXPipeline,
-    CogView4ControlPipeline,
     CogView4Pipeline,
 )
 
 TVideoPipeline = CogVideoXPipeline | CogVideoXImageToVideoPipeline
 TPipeline = CogView4Pipeline | TVideoPipeline
-CogviewPipline = CogView4Pipeline | CogView4ControlPipeline
 
 
 def _is_cogvideox1_0(pipeline: TVideoPipeline) -> bool:
@@ -105,7 +103,7 @@ def guess_resolution(
     height: int | None = None,
     width: int | None = None,
 ) -> tuple[int, int]:
-    if isinstance(pipeline, CogviewPipline):
+    if isinstance(pipeline, CogView4Pipeline):
         return _guess_cogview_resolution(pipeline, height=height, width=width)
     if isinstance(pipeline, TVideoPipeline):
         return _guess_cogvideox_resolution(pipeline, height=height, width=width)

src/cogkit/api/services/image_generation.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import os
+
 import torch
 
 from cogkit.api.logging import get_logger
@@ -26,12 +27,14 @@ def __init__(self, settings: APISettings) -> None:
             before_generation(cogview4_pl, settings.offload_type)
             self._models["cogview-4"] = cogview4_pl
 
+        ### Check if loaded models are supported
         for model in self._models.keys():
             if model not in settings._supported_models:
                 raise ValueError(
                     f"Registered model {model} not in supported list: {settings._supported_models}"
                 )
 
+        ### Check if all supported models are loaded
         for model in settings._supported_models:
             if model not in self._models:
                 _logger.warning(f"Model {model} not loaded")
@@ -54,6 +57,7 @@ def generate(
             raise ValueError(f"Model {model} not loaded")
         width, height = list(map(int, size.split("x")))
 
+        # TODO: Refactor this to switch by LoRA endpoint API
         if lora_path is not None:
             adapter_name = os.path.basename(lora_path)
             _logger.info(f"Loaded LORA weights from {adapter_name}")
@@ -64,12 +68,13 @@ def generate(
 
         output = generate_image(
             prompt=prompt,
+            pipeline=self._models[model],
+            num_images_per_prompt=num_images,
+            output_type="np",
             height=height,
             width=width,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
-            num_images_per_prompt=num_images,
-            output_type="np",
         )
 
         image_lst = self.postprocess(output)
@@ -79,7 +84,6 @@ def is_valid_model(self, model: str) -> bool:
         return model in self._models
 
     def postprocess(self, image_np: np.ndarray) -> list[np.ndarray]:
-        image_np = (image_np * 255).round().astype("uint8")
         image_lst = np.split(image_np, image_np.shape[0], axis=0)
         image_lst = [img.squeeze(0) for img in image_lst]
         return image_lst

src/cogkit/api/settings.py

@@ -19,5 +19,3 @@ class APISettings(BaseSettings):
     # cogview-4 related settings
     cogview4_path: str | None = None
     cogview4_transformer_path: str | None = None
-    openai_api_key: str | None = None
-    lora_dir: str | None = None

src/cogkit/cli/inference.py

@@ -31,18 +31,10 @@
     type=click.Path(dir_okay=False, writable=True),
     help="the path to save the generated image/video. If not provided, the generated image/video will be saved to 'output.png/mp4'.",
 )
-@click.option(
-    "--task",
-    type=click.Choice(
-        choices=[mode.value for mode in GenerationMode],
-        case_sensitive=False,
-    ),
-    help="the generation task",
-)
 @click.option(
     "--image_file",
     type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True),
-    help="the image to guide the video generation (for i2v or ct2i generation task)",
+    help="the image to guide the image/video generation (for i2i/i2v generation task)",
 )
 @click.option(
     "--dtype",
@@ -158,10 +150,7 @@ def inference(
         _logger.info("Saving the generated video to path '%s'.", os.fspath(output_file))
         export_to_video(output[0], output_file, fps=fps)
 
-    elif task in (
-        GenerationMode.TextToImage,
-        GenerationMode.CtrlTextToImage,
-    ):
+    elif task in (GenerationMode.TextToImage,):
         batched_images = generate_image(
             prompt=prompt,
             pipeline=pipeline,

src/cogkit/types/generation_mode.py

@@ -8,4 +8,3 @@ class GenerationMode(enum.Enum):
     TextToVideo = "t2v"
     ImageToVideo = "i2v"
     TextToImage = "t2i"
-    CtrlTextToImage = "ct2i"

src/cogkit/utils/__init__.py

@@ -6,7 +6,6 @@
 from cogkit.utils.lora import load_lora_checkpoint, unload_lora_checkpoint
 from cogkit.utils.misc import guess_generation_mode
 from cogkit.utils.path import mkdir, resolve_path
-from cogkit.utils.prompt import convert_prompt
 from cogkit.utils.random import rand_generator
 from cogkit.utils.load import load_pipeline
 
@@ -20,5 +19,4 @@
     "resolve_path",
     "rand_generator",
     "load_pipeline",
-    "convert_prompt",
 ]

src/cogkit/utils/misc.py

@@ -3,7 +3,7 @@
 
 from diffusers import DiffusionPipeline
 from PIL import Image
-from pathlib import Path
+
 from cogkit.logging import get_logger
 from cogkit.types import GenerationMode
 
@@ -14,7 +14,6 @@
     "CogView4Pipeline",
     "CogVideoXPipeline",
     "CogVideoXImageToVideoPipeline",
-    "CogView4ControlPipeline",
 )
 
 
@@ -37,23 +36,6 @@ def _check_text_to_image_params(
         )
 
 
-def _check_control_text_to_image_params(
-    pl_cls_name: str,
-    generation_mode: GenerationMode | None,
-    image: str | Path | None,
-) -> None:
-    if generation_mode is not None and generation_mode != GenerationMode.CtrlTextToImage:
-        _logger.warning(
-            "The pipeline `%s` does not support `%s` task. Will try the `%s` task.",
-            pl_cls_name,
-            generation_mode.value,
-            GenerationMode.CtrlTextToImage,
-        )
-    if image is not None:
-        err_msg = f"Image input is required in the image2video pipeline. Please provide a regular image file (image_file = {image})."
-        raise ValueError(err_msg)
-
-
 def _check_image_to_video_params(
     pl_cls_name: str,
     generation_mode: GenerationMode | None,
@@ -84,7 +66,7 @@ def guess_generation_mode(
     if generation_mode is not None:
         generation_mode = GenerationMode(generation_mode)
 
-    if pl_cls_name == "CogView4Pipeline":
+    if pl_cls_name.startswith("CogView"):
         # TextToImage
         _check_text_to_image_params(pl_cls_name, generation_mode, image)
         return GenerationMode.TextToImage
@@ -93,11 +75,6 @@ def guess_generation_mode(
         _check_image_to_video_params(pl_cls_name, generation_mode, image)
         return GenerationMode.ImageToVideo
 
-    if pl_cls_name == "CogView4ControlPipeline":
-        # Control TextToImage
-        _check_control_text_to_image_params(pl_cls_name, generation_mode, image)
-        return GenerationMode.CtrlTextToImage
-
     if image is not None:
         _logger.warning(
             "Pipeline `%s` does not support image input. Will ignore the image file.",