convaiinnovations/flux-test-time-training

1import torch
2from modeling_physics_rl import PhysicsModel, Config
3import os
4import sys
5
6def interactive_session():
7    print("\n============================================================")
8    print(" 🧪 FLUX TTT INFERENCE LAB (Pre-Trained)")
9    print("Commands:")
10    print("   - Type your question")
11    print("   - Type 'exit' to quit")
12    print("============================================================\n")
13
14    # 1. Load Model
15    print("🧠 Initializing Physics Model...")
16    model = PhysicsModel()
17    
18    # Force GPU if available
19    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
20    print(f"   🚀 Using Device: {device}")
21    model.to(device)
22    # Ensure inner LLM is also on device
23    model.llm.to(device)
24    
25    # 2. Load Trained TTT Weights
26    controller_path = "final_physics_controller.pt"
27    adapters_path = "final_flux_adapters.pt"
28    
29    try:
30        # Load Controller
31        if os.path.exists(controller_path):
32            print(f"   📂 Loading Controller: {controller_path}")
33            model.controller.load_state_dict(torch.load(controller_path, map_location=device))
34        else:
35            print(f"   ⚠️ Warning: Controller weights not found at {controller_path}")
36
37        # Load Flux Adapters
38        if os.path.exists(adapters_path):
39            print(f"   📂 Loading Flux Adapters: {adapters_path}")
40            states = torch.load(adapters_path, map_location=device)
41            # Handle list vs ModuleList vs simple state dict
42            if isinstance(states, list):
43                for layer, state in zip(model.flux_layers, states):
44                    layer.load_state_dict(state)
45            else:
46                model.flux_layers.load_state_dict(states)
47        else:
48            print(f"   ⚠️ Warning: Adapter weights not found at {adapters_path}")
49            
50    except Exception as e:
51        print(f"   ❌ Error loading weights: {e}")
52        print("   ⚠️ Proceeding with random/base weights...")
53
54    print("   ✅ Ready for Inference!\n")
55    
56    # 3. Interactive Loop
57    model.eval()
58    
59    while True:
60        try:
61            user_input = input("USER: ")
62            if user_input.lower() in ["exit", "quit"]:
63                break
64            
65            if not user_input.strip():
66                continue
67                
68            # Format prompt EXACTLY like training (System Prompt + Chat)
69            full_prompt = f"{Config.SYSTEM_PROMPT}\nUser: {user_input}\nModel:"
70            
71            inputs = model.tokenizer(full_prompt, return_tensors="pt").to(device)
72            
73            with torch.no_grad():
74                # 1. Predict Modulation
75                h_init = model.get_embeddings(inputs.input_ids).to(Config.DTYPE)
76                modulation = model.controller(h_init)
77                model.set_active_modulation(modulation)
78                
79                # 2. Generate Response
80                out_ids = model.llm.generate(
81                    **inputs,
82                    max_new_tokens=128,
83                    do_sample=True,
84                    temperature=0.6,    # Match Training (0.6)
85                    top_p=0.9,          # Match Training (0.9)
86                    repetition_penalty=1.2, # Match Training (1.2)
87                    pad_token_id=model.tokenizer.eos_token_id
88                )
89                
90                model.clear_modulation()
91                
92            response = model.tokenizer.decode(out_ids[0], skip_special_tokens=True)
93            
94            # Clean up response to show only the model's part
95            if "Model:" in response:
96                response = response.split("Model:")[-1].strip()
97            # Fallback cleanup just in case
98            elif response.startswith(full_prompt):
99                response = response[len(full_prompt):].strip()
100                
101            print(f"MODEL: {response}")
102            print(f"   [Modulation Norm: {torch.norm(modulation).item():.2f}]")
103            print("")
104            
105        except KeyboardInterrupt:
106            break
107        except Exception as e:
108            print(f"Error: {e}")
109
110if __name__ == "__main__":
111    interactive_session()
112