Update main.py

main.py

@@ -18,8 +18,14 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 # Load HF token from environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# Dictionary to store loaded pipelines
+loaded_pipelines = {}
+
 # Function to load pipeline dynamically
 def load_pipeline(model_name: str):
+    if model_name in loaded_pipelines:
+        return loaded_pipelines[model_name]
+
     if model_name == "Fluently XL Final":
         pipe = StableDiffusionXLPipeline.from_single_file(
             hf_hub_download(repo_id="fluently/Fluently-XL-Final", filename="FluentlyXL-Final.safetensors", token=HF_TOKEN),
@@ -65,6 +71,7 @@ def load_pipeline(model_name: str):
         raise ValueError(f"Unknown model: {model_name}")
     
     pipe.to(device)
+    loaded_pipelines[model_name] = pipe
     return pipe
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: