Updatr

app.py

@@ -22,25 +22,13 @@ MAX_IMAGE_SIZE = 2048
 # Global variables to store loaded components
 model = None
 scheduler = None
-text_encoder = None
-tokenizer = None
 decode_func = None
-null_cap_feat = None
-null_cap_mask = None
 config = None
 
 
 def load_model_components(device: str = "cuda"):
     """Load all model components once at startup"""
-    global \
-        model, \
-        scheduler, \
-        text_encoder, \
-        tokenizer, \
-        decode_func, \
-        null_cap_feat, \
-        null_cap_mask, \
-        config
+    global model, scheduler, decode_func, config
 
     try:
         # Load configuration
@@ -74,26 +62,6 @@ def load_model_components(device: str = "cuda"):
         else:
             raise ValueError("Unsupported VAE type")
 
-        print("Loading text encoder...")
-        # Load text encoder
-        text_encoder, tokenizer = load_text_encoder(
-            text_encoder_dir=model_config.text_encoder_dir,
-            device=device,
-            weight_dtype=torch.bfloat16,
-        )
-
-        print("Encoding null caption...")
-        # Get null caption features
-        null_cap_feat, null_cap_mask = encode_prompt(
-            tokenizer,
-            text_encoder,
-            device,
-            torch.bfloat16,
-            [""],
-            model_config.use_last_hidden_state,
-            max_seq_length=model_config.max_seq_length,
-        )
-
         print("Loading main model...")
         # Load main model
         model = instantiate_from_config(model_config.network).to(
@@ -129,6 +97,8 @@ def generate_image(
 ):
     """Generate image from text prompt"""
     try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {device}")
         # Validate inputs
         if not prompt or len(prompt.strip()) == 0:
             raise ValueError("Please enter a valid prompt")
@@ -136,9 +106,6 @@ def generate_image(
         if model is None or scheduler is None:
             raise RuntimeError("Model components not loaded. Please check the setup.")
 
-        if device == "cuda":
-            model.set_attn_implementation("flash_attention_2")
-
         # Validate dimensions
         if (
             width < 256
@@ -173,7 +140,14 @@ def generate_image(
             generator=generator,
         )
 
-        progress(0.1, desc="Encoding prompt...")
+        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,
+        )
 
         # Encode prompt
         cap_features, cap_mask = encode_prompt(
@@ -186,6 +160,17 @@ def generate_image(
             max_seq_length=config.model.max_seq_length,
         )
 
+        # Encode null caption for CFG
+        null_cap_feat, null_cap_mask = encode_prompt(
+            tokenizer,
+            text_encoder,
+            device,
+            dtype,
+            [""],
+            config.model.use_last_hidden_state,
+            max_seq_length=config.model.max_seq_length,
+        )
+
         cur_max_seq_len = cap_mask.sum(dim=-1).max()
         y = cap_features[:, :cur_max_seq_len]
 

tim/models/utils/text_encoders.py

@@ -11,7 +11,7 @@ def load_text_encoder(text_encoder_dir, device, weight_dtype):
         tokenizer.padding_side = "right"
         text_encoder = Gemma3ForCausalLM.from_pretrained(
             text_encoder_dir,
-            attn_implementation="sdpa",
+            attn_implementation="flash_attention_2",
             device_map="cpu",
             dtype=weight_dtype,
         ).model