Update main.py

main.py

@@ -9,6 +9,7 @@ import numpy as np
 import random
 from PIL import Image
 import io
+import os
 
 app = FastAPI()
 
@@ -16,91 +17,71 @@ MAX_SEED = np.iinfo(np.int32).max
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# 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
+# Load HF token from environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# Load pipelines
+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)
+
 
 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(...),
@@ -125,10 +106,8 @@ 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 = model_pipeline(
+        images = pipe_xl_final(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -139,7 +118,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently Anime":
-        images = model_pipeline(
+        images = pipe_anime(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -150,7 +129,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently Epic":
-        images = model_pipeline(
+        images = pipe_epic(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -161,7 +140,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently XL v4":
-        images = model_pipeline(
+        images = pipe_xl(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -172,7 +151,7 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently XL v3 Lightning":
-        images = model_pipeline(
+        images = pipe_xl_lightning(
             prompt=prompt,
             negative_prompt=negative_prompt,
             width=width,
@@ -183,8 +162,8 @@ async def generate(
             output_type="pil",
         ).images
     elif model == "Fluently v4 inpaint" or model == "Fluently XL v3 inpaint":
-        blurred_mask = model_pipeline.mask_processor.blur(mask_image_pil, blur_factor=blur_factor)
-        images = model_pipeline(
+        blurred_mask = pipe_inpaint.mask_processor.blur(mask_image_pil, blur_factor=blur_factor)
+        images = pipe_inpaint(
             prompt=prompt,
             image=inpaint_image_pil,
             mask_image=blurred_mask,
@@ -198,10 +177,6 @@ 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')