Update main.py

main.py

@@ -16,69 +16,91 @@ MAX_SEED = np.iinfo(np.int32).max
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Load pipelines
-HF_TOKEN = os.getenv("HF_TOKEN") # Replace with your actual token
-pipe_xl_final = StableDiffusionXLPipeline.from_single_file(
-    hf_hub_download(repo_id="fluently/Fluently-XL-Final", filename="FluentlyXL-Final.safetensors", token=HF_TOKEN),
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipe_xl_final.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_xl_final.scheduler.config)
-pipe_xl_final.to(device)
-
-pipe_anime = StableDiffusionPipeline.from_pretrained(
-    "fluently/Fluently-anime",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipe_anime.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_anime.scheduler.config)
-pipe_anime.to(device)
-
-pipe_epic = StableDiffusionPipeline.from_pretrained(
-    "fluently/Fluently-epic",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipe_epic.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_epic.scheduler.config)
-pipe_epic.to(device)
-
-pipe_xl_inpaint = StableDiffusionXLInpaintPipeline.from_single_file(
-    "https://huggingface.co/fluently/Fluently-XL-v3-inpainting/blob/main/FluentlyXL-v3-inpainting.safetensors",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipe_xl_inpaint.to(device)
-
-pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
-    "fluently/Fluently-v4-inpainting",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipe_inpaint.to(device)
-
-pipe_xl = StableDiffusionXLPipeline.from_pretrained(
-    "fluently/Fluently-XL-v4",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipe_xl.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_xl.scheduler.config)
-pipe_xl.to(device)
-
-pipe_xl_lightning = StableDiffusionXLPipeline.from_pretrained(
-    "fluently/Fluently-XL-v3-lightning",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-)
-pipe_xl_lightning.scheduler = DPMSolverSinglestepScheduler.from_config(pipe_xl_lightning.scheduler.config, use_karras_sigmas=False, timestep_spacing="trailing", lower_order_final=True)
-pipe_xl_lightning.to(device)
-
+# Set directory for model storage
+MODEL_DIR = "/data"
+
+# Ensure model directory exists
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+# Download models to local directory
+HF_TOKEN = os.getenv("HF_TOKEN")  # Replace with your actual token
+def download_model(repo_id, filename=None, model_dir=MODEL_DIR, token=HF_TOKEN):
+    if filename:
+        return hf_hub_download(repo_id=repo_id, filename=filename, local_dir=model_dir, token=token)
+    return hf_hub_download(repo_id=repo_id, local_dir=model_dir, token=token)
+
+# Paths for models
+paths = {
+    "Fluently XL Final": download_model("fluently/Fluently-XL-Final", "FluentlyXL-Final.safetensors"),
+    "Fluently Anime": download_model("fluently/Fluently-anime"),
+    "Fluently Epic": download_model("fluently/Fluently-epic"),
+    "Fluently XL v4": download_model("fluently/Fluently-XL-v4"),
+    "Fluently XL v3 Lightning": download_model("fluently/Fluently-XL-v3-lightning"),
+    "Fluently v4 inpaint": download_model("fluently/Fluently-v4-inpainting"),
+    "Fluently XL v3 inpaint": download_model("fluently/Fluently-XL-v3-inpainting", "FluentlyXL-v3-inpainting.safetensors"),
+}
+
+# Function to load model dynamically
+def load_model(model_name):
+    if model_name == "Fluently XL Final":
+        model = StableDiffusionXLPipeline.from_single_file(
+            paths[model_name],
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config)
+    elif model_name == "Fluently Anime":
+        model = StableDiffusionPipeline.from_pretrained(
+            paths[model_name],
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config)
+    elif model_name == "Fluently Epic":
+        model = StableDiffusionPipeline.from_pretrained(
+            paths[model_name],
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config)
+    elif model_name == "Fluently XL v4":
+        model = StableDiffusionXLPipeline.from_pretrained(
+            paths[model_name],
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config)
+    elif model_name == "Fluently XL v3 Lightning":
+        model = StableDiffusionXLPipeline.from_pretrained(
+            paths[model_name],
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        model.scheduler = DPMSolverSinglestepScheduler.from_config(model.scheduler.config, use_karras_sigmas=False, timestep_spacing="trailing", lower_order_final=True)
+    elif model_name in ["Fluently v4 inpaint", "Fluently XL v3 inpaint"]:
+        if model_name == "Fluently v4 inpaint":
+            model = StableDiffusionInpaintPipeline.from_pretrained(
+                paths[model_name],
+                torch_dtype=torch.float16,
+                use_safetensors=True,
+            )
+        else:
+            model = StableDiffusionXLInpaintPipeline.from_single_file(
+                paths[model_name],
+                torch_dtype=torch.float16,
+                use_safetensors=True,
+            )
+    else:
+        raise ValueError(f"Model {model_name} not found")
+    
+    model.to(device)
+    return model
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-
 @app.post("/generate")
 async def generate(
     model: str = Form(...),
@@ -103,8 +125,10 @@ async def generate(
     inpaint_image_pil = Image.open(io.BytesIO(await inpaint_image.read())) if inpaint_image else None
     mask_image_pil = Image.open(io.BytesIO(await mask_image.read())) if mask_image else None
 
+    model_pipeline = load_model(model)
+
     if model == "Fluently XL Final":
-        images = pipe_xl_final(
+        images = model_pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -115,7 +139,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently Anime":
-        images = pipe_anime(
+        images = model_pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -126,7 +150,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently Epic":
-        images = pipe_epic(
+        images = model_pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -137,7 +161,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently XL v4":
-        images = pipe_xl(
+        images = model_pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -148,7 +172,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently XL v3 Lightning":
-        images = pipe_xl_lightning(
+        images = model_pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -159,8 +183,8 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently v4 inpaint" or model == "Fluently XL v3 inpaint":
-        blurred_mask = pipe_inpaint.mask_processor.blur(mask_image_pil, blur_factor=blur_factor)
-        images = pipe_inpaint(
+        blurred_mask = model_pipeline.mask_processor.blur(mask_image_pil, blur_factor=blur_factor)
+        images = model_pipeline(
             prompt=prompt,
             image=inpaint_image_pil,
             mask_image=blurred_mask,
@@ -174,6 +198,10 @@ async def generate(
             output_type="pil",
         ).images
 
+    # Unload the model from the device
+    model_pipeline.to("cpu")
+    torch.cuda.empty_cache()
+
     img = images[0]
     img_byte_arr = io.BytesIO()
     img.save(img_byte_arr, format='PNG')