Renaming

app.py

@@ -8,11 +8,11 @@ import google.generativeai as genai
 
 import spaces
 import torch
-from pikigen import PikigenPipeline
+from f_lite import FLitePipeline
 
 # Trick required because it is not a native diffusers model
 from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMPORTABLE_CLASSES
-LOADABLE_CLASSES["pikigen"] = LOADABLE_CLASSES["pikigen.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
+LOADABLE_CLASSES["f_lite"] = LOADABLE_CLASSES["f_lite.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
 ALL_IMPORTABLE_CLASSES["DiT"] = ["save_pretrained", "from_pretrained"]
 
 load_dotenv()
@@ -33,7 +33,7 @@ if torch.cuda.is_available():
 else:
     torch_dtype = torch.float32
 
-pipe = PikigenPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pipe = FLitePipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 # pipe.enable_model_cpu_offload()  # For less memory consumption
 pipe.to(device)
 pipe.vae.enable_slicing()

f_lite/__init__.py

@@ -0,0 +1,5 @@
+from .pipeline import FLitePipeline, FLitePipelineOutput, APGConfig
+from .model import DiT
+
+
+__all__ = ["FLitePipeline", "FLitePipelineOutput", "APGConfig", "DiT"]

{pikigen → f_lite}/pipeline.py

@@ -27,9 +27,9 @@ class APGConfig:
 
 
 @dataclass
-class PikigenPipelineOutput(BaseOutput):
+class FLitePipelineOutput(BaseOutput):
     """
-    Output class for PikigenPipeline pipeline.
+    Output class for FLitePipeline pipeline.
     Args:
         images (`List[PIL.Image.Image]` or `np.ndarray`)
             List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
@@ -39,9 +39,9 @@ class PikigenPipelineOutput(BaseOutput):
     images: Union[List[Image.Image], np.ndarray]
 
 
-class PikigenPipeline(DiffusionPipeline):
+class FLitePipeline(DiffusionPipeline):
     r"""
-    Pipeline for text-to-image generation using Pikigen model.
+    Pipeline for text-to-image generation using FLite model.
     This model inherits from [`DiffusionPipeline`].
     """
 
@@ -289,7 +289,7 @@ class PikigenPipeline(DiffusionPipeline):
         images = (images * 255).round().clamp(0, 255).to(torch.uint8).cpu()
         pil_images = [Image.fromarray(img.permute(1, 2, 0).numpy()) for img in images]
 
-        return PikigenPipelineOutput(
+        return FLitePipelineOutput(
             images=pil_images,
         )
 

pikigen/__init__.py

@@ -1,5 +0,0 @@
-from .pipeline import PikigenPipeline, PikigenPipelineOutput, APGConfig
-from .model import DiT
-
-
-__all__ = ["PikigenPipeline", "PikigenPipelineOutput", "APGConfig", "DiT"]