Fix loading problem

app.py

@@ -9,6 +9,7 @@ from pikigen import PikigenPipeline
 # Trick required because it is not a native diffusers model
 from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMPORTABLE_CLASSES
 LOADABLE_CLASSES.setdefault("pikigen", {}).setdefault("DiT", []).extend(["save_pretrained", "from_pretrained"])
+LOADABLE_CLASSES["pikigen.model"] = LOADABLE_CLASSES["pikigen"]
 ALL_IMPORTABLE_CLASSES.setdefault("DiT", []).extend(["save_pretrained", "from_pretrained"])
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -20,7 +21,7 @@ else:
     torch_dtype = torch.float32
 
 pipe = PikigenPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe.enable_model_cpu_offload()  # For less memory consumption
+# pipe.enable_model_cpu_offload()  # For less memory consumption
 pipe.vae.enable_slicing()
 pipe.vae.enable_tiling()