Update main.py

main.py

@@ -3,80 +3,75 @@ from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import StreamingResponse
 import torch
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, DPMSolverSinglestepScheduler
-from diffusers.pipelines import StableDiffusionInpaintPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.pipelines import StableDiffusionInpaintPipeline
 from huggingface_hub import hf_hub_download
 import numpy as np
 import random
 from PIL import Image
 import io
-import os
 
 app = FastAPI()
 
 MAX_SEED = np.iinfo(np.int32).max
-
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # 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_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)
-
+# Function to load pipeline dynamically
+def load_pipeline(model_name: str):
+    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),
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif model_name == "Fluently Anime":
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "fluently/Fluently-anime",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif model_name == "Fluently Epic":
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "fluently/Fluently-epic",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif model_name == "Fluently XL v4":
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            "fluently/Fluently-XL-v4",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    elif model_name == "Fluently XL v3 Lightning":
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            "fluently/Fluently-XL-v3-lightning",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=False, timestep_spacing="trailing", lower_order_final=True)
+    elif model_name == "Fluently v4 inpaint":
+        pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            "fluently/Fluently-v4-inpainting",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+    
+    pipe.to(device)
+    return pipe
 
 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(...),
@@ -101,73 +96,31 @@ 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
 
-    if model == "Fluently XL Final":
-        images = pipe_xl_final(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale,
-            num_inference_steps=25,
-            num_images_per_prompt=1,
-            output_type="pil",
-        ).images
-    elif model == "Fluently Anime":
-        images = pipe_anime(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale,
-            num_inference_steps=30,
-            num_images_per_prompt=1,
-            output_type="pil",
-        ).images
-    elif model == "Fluently Epic":
-        images = pipe_epic(
+    pipe = load_pipeline(model)
+
+    if model in ["Fluently v4 inpaint"]:
+        blurred_mask = pipe.mask_processor.blur(mask_image_pil, blur_factor=blur_factor)
+        images = pipe(
             prompt=prompt,
+            image=inpaint_image_pil,
+            mask_image=blurred_mask,
             negative_prompt=negative_prompt,
             width=width,
             height=height,
             guidance_scale=guidance_scale,
             num_inference_steps=30,
+            strength=strength,
             num_images_per_prompt=1,
             output_type="pil",
         ).images
-    elif model == "Fluently XL v4":
-        images = pipe_xl(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale,
-            num_inference_steps=25,
-            num_images_per_prompt=1,
-            output_type="pil",
-        ).images
-    elif model == "Fluently XL v3 Lightning":
-        images = pipe_xl_lightning(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            width=width,
-            height=height,
-            guidance_scale=2,
-            num_inference_steps=5,
-            num_images_per_prompt=1,
-            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(
+    else:
+        images = pipe(
             prompt=prompt,
-            image=inpaint_image_pil,
-            mask_image=blurred_mask,
             negative_prompt=negative_prompt,
             width=width,
             height=height,
             guidance_scale=guidance_scale,
-            num_inference_steps=30,
-            strength=strength,
+            num_inference_steps=25 if model == "Fluently XL Final" else 30,
             num_images_per_prompt=1,
             output_type="pil",
         ).images
@@ -179,7 +132,6 @@ async def generate(
 
     return StreamingResponse(img_byte_arr, media_type="image/png")
 
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)