Update src/pipeline.py

src/pipeline.py

@@ -49,14 +49,6 @@ def load_quanto_text_encoder_2(text_repo_path):
 
 def load_pipeline() -> Pipeline:
 
-    try:
-        text_repo_path = os.path.join(HF_HUB_CACHE, "models--RichardWilliam--XULF_T5_bf16/snapshots/63a3d9ef7b586655600ac9bd4e4747d038237761")
-        text_encoder_2 = load_quanto_text_encoder_2(text_repo_path=text_repo_path)
-    except:
-        text_encoder_2 =  T5EncoderModel.from_pretrained("RichardWilliam/XULF_T5_bf16", 
-                    revision = "63a3d9ef7b586655600ac9bd4e4747d038237761", 
-                    torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
-
     origin_vae = AutoencoderTiny.from_pretrained("RichardWilliam/XULF_Vae",
                     revision="3ee225c539465c27adadec45c6e8af50a7397b7d", 
                     torch_dtype=torch.bfloat16)
@@ -69,15 +61,20 @@ def load_pipeline() -> Pipeline:
     transformer = origin_trans
 
     pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT, 
-                        revision=REVISION, 
-                        vae=origin_vae, 
+                        revision=REVISION,
                         transformer=transformer, 
-                        text_encoder_2=text_encoder_2, 
                         torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
 
+    try:
+        # pipeline.text_encoder_v2 = load_quanto_text_encoder_2(text_repo_path=None)
+        pipeline.enable_cuda_graph(type="max-autotune")
+        pipeline.text_encoder_v2 = load_quanto_text_encoder_2(text_repo_path=None)
+    except:
+        print("Something wrong here")
+
     for __ in range(3):
-        pipeline(prompt="sweet, subordinative, gender, mormyre, arteriolosclerosis, positivism, Antiochianism, palmerite", 
+        pipeline(prompt="schoenobatist, halisteresis, chronometric, hallucinative", 
                         width=1024, 
                         height=1024, 
                         guidance_scale=0.0, 
@@ -89,6 +86,7 @@ def load_pipeline() -> Pipeline:
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
 
     reset_cache()
+    torch.cuda.empty_cache()
 
     generator = Generator(pipeline.device).manual_seed(request.seed)