clean up

server/config.py

@@ -16,8 +16,6 @@ class Args(BaseModel):
     pipeline: str
     ssl_certfile: str | None
     ssl_keyfile: str | None
-    sfast: bool
-    onediff: bool = False
     compel: bool = False
     debug: bool = False
     pruna: bool = False
@@ -112,18 +110,6 @@ parser.add_argument(
     default=False,
     help="Compel",
 )
-parser.add_argument(
-    "--sfast",
-    action="store_true",
-    default=False,
-    help="Enable Stable Fast",
-)
-parser.add_argument(
-    "--onediff",
-    action="store_true",
-    default=False,
-    help="Enable OneDiff",
-)
 parser.add_argument(
     "--pruna",
     action="store_true",

server/main.py

@@ -5,7 +5,6 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from fastapi import Request
 import markdown2
-from pipelines.utils.safety_checker import SafetyChecker
 from PIL import Image
 import logging
 from config import config, Args
@@ -53,8 +52,10 @@ class App:
         self.pipeline = pipeline_instance
         self.app = FastAPI()
         self.conn_manager = ConnectionManager()
-        self.safety_checker: SafetyChecker | None = None
+        self.safety_checker = None
         if self.args.safety_checker:
+            from pipelines.utils.safety_checker import SafetyChecker
+
             self.safety_checker = SafetyChecker(device=device.type)
         self.init_app()
 
@@ -80,9 +81,13 @@ class App:
                 # Handle websocket disconnection event
                 code = disconnect_error.code
                 if code == 1006:  # ABNORMAL_CLOSURE
-                    logging.info(f"WebSocket abnormally closed for user {user_id}: Connection was closed without a proper close handshake")
+                    logging.info(
+                        f"WebSocket abnormally closed for user {user_id}: Connection was closed without a proper close handshake"
+                    )
                 else:
-                    logging.info(f"WebSocket disconnected for user {user_id} with code {code}: {disconnect_error.reason}")
+                    logging.info(
+                        f"WebSocket disconnected for user {user_id} with code {code}: {disconnect_error.reason}"
+                    )
             except RuntimeError as e:
                 if any(err in str(e) for err in ERROR_MESSAGES):
                     logging.info(f"WebSocket disconnected for user {user_id}: {e}")
@@ -268,10 +273,14 @@ class App:
                 # Handle websocket disconnection event
                 code = disconnect_error.code
                 if code == 1006:  # ABNORMAL_CLOSURE
-                    logging.info(f"WebSocket abnormally closed during streaming for user {user_id}: Connection was closed without a proper close handshake")
+                    logging.info(
+                        f"WebSocket abnormally closed during streaming for user {user_id}: Connection was closed without a proper close handshake"
+                    )
                 else:
-                    logging.info(f"WebSocket disconnected during streaming for user {user_id} with code {code}: {disconnect_error.reason}")
-                
+                    logging.info(
+                        f"WebSocket disconnected during streaming for user {user_id} with code {code}: {disconnect_error.reason}"
+                    )
+
                 # Clean disconnection without error response
                 await self.conn_manager.disconnect(user_id)
                 raise HTTPException(status_code=204, detail="Connection closed")

server/pipelines/IPcompositionHyperSD15.py

@@ -129,19 +129,6 @@ class Pipeline:
         # pipe.unet = oneflow_compile(pipe.unet, options=compile_options)
         # pipe.vae.decoder = oneflow_compile(pipe.vae.decoder, options=compile_options)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            # cofig.
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/IPcompositionHyperSDXL.py

@@ -130,18 +130,6 @@ class Pipeline:
         self.pipe.scheduler = TCDScheduler.from_config(self.pipe.scheduler.config)
         self.pipe.set_ip_adapter_scale([0.8])
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnet.py

@@ -179,28 +179,6 @@ class Pipeline:
             smash_config["compiler"] = "stable_fast"
             self.pipe = smash(model=self.pipe, smash_config=smash_config)
 
-        if args.sfast:
-            print("\nRunning sfast compile\n")
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
-        if args.onediff:
-            print("\nRunning onediff compile\n")
-            from onediff.infer_compiler import oneflow_compile
-
-            self.pipe.unet = oneflow_compile(self.pipe.unet)
-            self.pipe.vae.encoder = oneflow_compile(self.pipe.vae.encoder)
-            self.pipe.vae.decoder = oneflow_compile(self.pipe.vae.decoder)
-            self.pipe.controlnet = oneflow_compile(self.pipe.controlnet)
-
         self.canny_torch = SobelOperator(device=device)
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)

server/pipelines/controlnetDepthFlashSD.py

@@ -169,18 +169,6 @@ class Pipeline:
         )
         self.pipe.fuse_lora()
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetDepthHyperSD.py

@@ -169,18 +169,6 @@ class Pipeline:
 
         self.pipe.fuse_lora()
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetDepthHyperSDXL.py

@@ -168,18 +168,6 @@ class Pipeline:
 
         self.pipe.fuse_lora()
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetFlashSD.py

@@ -161,18 +161,6 @@ class Pipeline:
 
         self.canny_torch = SobelOperator(device=device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetFlashSDXL.py

@@ -162,18 +162,6 @@ class Pipeline:
 
         self.canny_torch = SobelOperator(device=device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetHyperSD.py

@@ -177,18 +177,6 @@ class Pipeline:
         self.pipe.fuse_lora()
         self.canny_torch = SobelOperator(device=device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetHyperSDXL.py

@@ -182,18 +182,6 @@ class Pipeline:
         self.pipe.fuse_lora()
         self.canny_torch = SobelOperator(device=device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetLoraSDXL-Lightning.py

@@ -202,18 +202,6 @@ class Pipeline:
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         if device.type != "mps":
             self.pipe.unet.to(memory_format=torch.channels_last)
 

server/pipelines/controlnetLoraSDXL.py

@@ -194,18 +194,6 @@ class Pipeline:
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype).to(device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         if device.type != "mps":
             self.pipe.unet.to(memory_format=torch.channels_last)
 

server/pipelines/controlnetMistoLineHyperSDXL.py

@@ -186,18 +186,6 @@ class Pipeline:
             "TheMistoAI/MistoLine", filename="MTEED.pth", subfolder="Anyline"
         ).to(device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            # config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device)
         if device.type != "mps":

server/pipelines/controlnetSDTurbo.py

@@ -174,28 +174,6 @@ class Pipeline:
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
             ).to(device)
 
-        if args.sfast:
-            print("\nRunning sfast compile\n")
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
-        if args.onediff:
-            print("\nRunning onediff compile\n")
-            from onediff.infer_compiler import oneflow_compile
-
-            self.pipe.unet = oneflow_compile(self.pipe.unet)
-            self.pipe.vae.encoder = oneflow_compile(self.pipe.vae.encoder)
-            self.pipe.vae.decoder = oneflow_compile(self.pipe.vae.decoder)
-            self.pipe.controlnet = oneflow_compile(self.pipe.controlnet)
-
         self.canny_torch = SobelOperator(device=device)
 
         self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)

server/pipelines/controlnetSDXLTurbo.py

@@ -180,18 +180,6 @@ class Pipeline:
         )
         self.canny_torch = SobelOperator(device=device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype).to(device)
         if device.type != "mps":

server/pipelines/controlnetSegmindVegaRT.py

@@ -188,18 +188,6 @@ class Pipeline:
             base_model, subfolder="scheduler"
         )
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype).to(device)
         if device.type != "mps":

server/pipelines/img2img.py

@@ -105,18 +105,6 @@ class Pipeline:
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
             ).to(device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
         if device.type != "mps":

server/pipelines/img2imgSDTurbo.py

@@ -102,7 +102,7 @@ class Pipeline:
         if args.pruna:
             # Create and smash your model
             smash_config = SmashConfig()
-            # smash_config["cacher"] = "deepcache"
+            smash_config["cacher"] = "deepcache"
             smash_config["compiler"] = "stable_fast"
             self.pipe = smash(model=self.pipe, smash_config=smash_config)
 

server/pipelines/img2imgSDXL-Lightning.py

@@ -143,18 +143,6 @@ class Pipeline:
             smash_config["compiler"] = "stable_fast"
             self.pipe = smash(model=self.pipe, smash_config=smash_config)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
         if device.type != "mps":

server/pipelines/img2imgSDXLTurbo.py

@@ -121,18 +121,6 @@ class Pipeline:
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
             ).to(device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         if device.type != "mps":
             self.pipe.unet.to(memory_format=torch.channels_last)
 

server/pipelines/img2imgSDXS512.py

@@ -102,18 +102,6 @@ class Pipeline:
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
             ).to(device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
         if device.type != "mps":

server/pipelines/img2imgSegmindVegaRT.py

@@ -121,17 +121,6 @@ class Pipeline:
         self.pipe.scheduler = LCMScheduler.from_pretrained(
             base_model, subfolder="scheduler"
         )
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
 
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)

server/pipelines/txt2img.py

@@ -94,18 +94,6 @@ class Pipeline:
             smash_config["compiler"] = "stable_fast"
             self.pipe = smash(model=self.pipe, smash_config=smash_config)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
         if device.type != "mps":

server/pipelines/txt2imgLora.py

@@ -119,18 +119,6 @@ class Pipeline:
 
             self.pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         if args.compel:
             self.compel_proc = Compel(
                 tokenizer=self.pipe.tokenizer,

server/pipelines/txt2imgLoraSDXL.py

@@ -108,18 +108,6 @@ class Pipeline:
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype).to(device)
 
-        if args.sfast:
-            from sfast.compilers.stable_diffusion_pipeline_compiler import (
-                compile,
-                CompilationConfig,
-            )
-
-            config = CompilationConfig.Default()
-            config.enable_xformers = True
-            config.enable_triton = True
-            config.enable_cuda_graph = True
-            self.pipe = compile(self.pipe, config=config)
-
         if device.type != "mps":
             self.pipe.unet.to(memory_format=torch.channels_last)
 

server/pipelines/utils/safety_checker.py

@@ -143,10 +143,10 @@ class SafetyChecker:
 
         self.device = device
         self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
-            "CompVis/stable-diffusion-safety-checker"
+            "CompVis/stable-diffusion-safety-checker",
         ).to(device)
         self.feature_extractor = CLIPFeatureExtractor.from_pretrained(
-            "openai/clip-vit-base-patch32"
+            "openai/clip-vit-base-patch32",
         )
 
     def __call__(

server/requirements.txt

@@ -20,7 +20,7 @@ setuptools
 mpmath
 controlnet-aux
 sentencepiece
-optimum-quanto # has to be optimum-quanto==0.2.5 for pruna int4
+optimum-quanto==0.2.5
 gguf
 types-Pillow
 mypy