soijds

app.py

@@ -24,11 +24,13 @@ model = None
 scheduler = None
 decode_func = None
 config = None
+text_encoder = None
+tokenizer = None
 
 
 def load_model_components(device: str = "cuda"):
     """Load all model components once at startup"""
-    global model, scheduler, decode_func, config
+    global model, scheduler, decode_func, config, text_encoder, tokenizer
 
     try:
         # Load configuration
@@ -62,6 +64,13 @@ def load_model_components(device: str = "cuda"):
         else:
             raise ValueError("Unsupported VAE type")
 
+        # Load text encoder
+        text_encoder, tokenizer = load_text_encoder(
+            text_encoder_dir=config.model.text_encoder_dir,
+            device=device,
+            weight_dtype=dtype,
+        )
+
         print("Loading main model...")
         # Load main model
         model = instantiate_from_config(model_config.network).to(
@@ -143,16 +152,13 @@ def generate_image(
         progress(0.1, desc="Loading text encoder...")
 
         # Load text encoder
-        text_encoder, tokenizer = load_text_encoder(
-            text_encoder_dir=config.model.text_encoder_dir,
-            device=device,
-            weight_dtype=dtype,
-        )
+        text_encoder.to(device)
+        text_encoder.set_attn_implementation("flash_attention_2")
 
         # Encode prompt
         cap_features, cap_mask = encode_prompt(
             tokenizer,
-            text_encoder,
+            text_encoder.model,
             device,
             dtype,
             [prompt],
@@ -163,7 +169,7 @@ def generate_image(
         # Encode null caption for CFG
         null_cap_feat, null_cap_mask = encode_prompt(
             tokenizer,
-            text_encoder,
+            text_encoder.model,
             device,
             dtype,
             [""],

tim/models/utils/text_encoders.py

@@ -11,10 +11,10 @@ def load_text_encoder(text_encoder_dir, device, weight_dtype):
         tokenizer.padding_side = "right"
         text_encoder = Gemma3ForCausalLM.from_pretrained(
             text_encoder_dir,
-            attn_implementation="flash_attention_2",
+            attn_implementation="sdpa",
             device_map="cpu",
             dtype=weight_dtype,
-        ).model
+        )
     elif "t5" in text_encoder_dir:
         text_encoder = T5EncoderModel.from_pretrained(
             text_encoder_dir,
@@ -28,7 +28,7 @@ def load_text_encoder(text_encoder_dir, device, weight_dtype):
     # for param in text_encoder.parameters():
     #     param.requires_grad = False
 
-    text_encoder = text_encoder.eval().to(device=device, dtype=weight_dtype)
+    text_encoder.model = text_encoder.model.eval().to(device=device, dtype=weight_dtype)
 
     return text_encoder, tokenizer