Update src/pipeline.py

src/pipeline.py

@@ -788,11 +788,9 @@ os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 torch._dynamo.config.suppress_errors = True
 
-Pipeline: TypeAlias = FluxPipeline
-
-CHECKPOINT = "black-forest-labs/FLUX.1-schnell"
-REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
-
+Pipeline = None
+ids = "slobers/Flux.1.Schnella"
+Revision = "e34d670e44cecbbc90e4962e7aada2ac5ce8b55b"
 
 def load_pipeline() -> Pipeline:
     text_encoder = CLIPTextModel.from_pretrained(CHECKPOINT, revision=REVISION, subfolder="text_encoder", local_files_only=True, torch_dtype=torch.bfloat16,)
@@ -826,11 +824,8 @@ def load_pipeline() -> Pipeline:
 
     return pipeline
 
+@torch.no_grad()
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-
     generator = Generator(pipeline.device).manual_seed(request.seed)
 
     return pipeline(
@@ -842,3 +837,4 @@ def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
         height=request.height,
         width=request.width,
     ).images[0]
+