Flux Pipeline Discard

Author: FredyRivera-dev

examples/server-async/Pipelines.py

@@ -1,5 +1,4 @@
 from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import StableDiffusion3Pipeline
-from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
 import torch
 import os
 import logging
@@ -20,7 +19,6 @@ class TextToImageInput(BaseModel):
 class PresetModels:
     SD3: List[str] = field(default_factory=lambda: ['stabilityai/stable-diffusion-3-medium'])
     SD3_5: List[str] = field(default_factory=lambda: ['stabilityai/stable-diffusion-3.5-large', 'stabilityai/stable-diffusion-3.5-large-turbo', 'stabilityai/stable-diffusion-3.5-medium'])
-    Flux: List[str] = field(default_factory=lambda: ['black-forest-labs/FLUX.1-dev', 'black-forest-labs/FLUX.1-schnell'])
 
 class TextToImagePipelineSD3:
     def __init__(self, model_path: str | None = None):
@@ -48,37 +46,6 @@ def start(self):
         else:
             raise Exception("No CUDA or MPS device available")
 
-class TextToImagePipelineFlux:
-    def __init__(self, model_path: str | None = None, low_vram: bool = False):
-        self.model_path = model_path or os.getenv("MODEL_PATH")
-        self.pipeline: FluxPipeline | None = None
-        self.device: str | None = None
-        self.low_vram = low_vram
-
-    def start(self):
-        if torch.cuda.is_available():
-            model_path = self.model_path or "black-forest-labs/FLUX.1-schnell"
-            logger.info("Loading CUDA")
-            self.device = "cuda" 
-            self.pipeline = FluxPipeline.from_pretrained(
-                model_path,
-                torch_dtype=torch.bfloat16,
-            ).to(device=self.device)
-            if self.low_vram:
-                self.pipeline.enable_model_cpu_offload()
-            else:
-                pass
-        elif torch.backends.mps.is_available():
-            model_path = self.model_path or "black-forest-labs/FLUX.1-schnell"
-            logger.info("Loading MPS for Mac M Series")
-            self.device = "mps"
-            self.pipeline = FluxPipeline.from_pretrained(
-                model_path,
-                torch_dtype=torch.bfloat16,
-            ).to(device=self.device)
-        else:
-            raise Exception("No CUDA or MPS device available")
-
 class ModelPipelineInitializer:
     def __init__(self, model: str = '', type_models: str = 't2im'):
         self.model = model
@@ -99,15 +66,11 @@ def initialize_pipeline(self):
             self.model_type = "SD3"
         elif self.model in preset_models.SD3_5:
             self.model_type = "SD3_5"
-        elif self.model in preset_models.Flux:
-            self.model_type = "Flux"
 
         # Create appropriate pipeline based on model type and type_models
         if self.type_models == 't2im':
             if self.model_type in ["SD3", "SD3_5"]:
                 self.pipeline = TextToImagePipelineSD3(self.model)
-            elif self.model_type == "Flux":
-                self.pipeline = TextToImagePipelineFlux(self.model)
             else:
                 raise ValueError(f"Model type {self.model_type} not supported for text-to-image")
         elif self.type_models == 't2v':