Increase duration

app.py

@@ -8,9 +8,8 @@ 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"])
+LOADABLE_CLASSES["pikigen"] = LOADABLE_CLASSES["pikigen.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
+ALL_IMPORTABLE_CLASSES["DiT"] = ["save_pretrained", "from_pretrained"]
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Freepik/Pikigen-test"
@@ -22,6 +21,7 @@ else:
 
 pipe = PikigenPipeline.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()
 pipe.vae.enable_tiling()
 
@@ -29,7 +29,7 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
-@spaces.GPU
+@spaces.GPU(duration=500)
 def infer(
     prompt,
     negative_prompt,