Make model global

app.py

@@ -73,7 +73,7 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0):
             noised[idx] = val
     return noised
 
-def generate_diffusion_text(model, input_ids, answer_start):
+def generate_diffusion_text(input_ids, answer_start):
     with torch.no_grad():
         input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device)
         logits = model(input_ids=input_tensor)["logits"]
@@ -83,7 +83,7 @@ def generate_diffusion_text(model, input_ids, answer_start):
     return input_ids[:answer_start] + sampled[answer_start:]
 
 # --- Inference Wrapper ---
-def diffusion_chat(question, eot_weight, max_it, sharpness, model):
+def diffusion_chat(question, eot_weight, max_it, sharpness):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
@@ -106,7 +106,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, model):
     last_tokens = []
 
     for i in range(max_it):
-        generated_tokens = generate_diffusion_text(model, current_tokens, answer_start)
+        generated_tokens = generate_diffusion_text(current_tokens, answer_start)
         current_tokens = generated_tokens
 
         decoded_ids = current_tokens[answer_start:]
@@ -144,7 +144,10 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, model):
     yield f"<b>Final Output (after {i+1} iterations):</b><br>" + final_output
 
 # --- Gradio Interface ---
-model_state = gr.State(value=load_model())  # this just stores the object
+
+print("Loading model...")
+model = load_model()
+print("✅ Model loaded.")
 
 demo = gr.Interface(
     fn=diffusion_chat,
@@ -152,8 +155,7 @@ demo = gr.Interface(
         gr.Textbox(label="User Question", lines=2, placeholder="What do you know about the city of New York?"),
         gr.Slider(0, 1, value=0.4, step=0.05, label="↓ = longer answers (EOT weight)"),
         gr.Slider(1, 512, value=64, step=1, label="↑ = more iterations"),
-        gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)"),
-        model_state
+        gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)")
     ],
     outputs=gr.HTML(label="Diffusion Output"),
     title="Diffusion Language Model Chat",