app.py

# dream_app.py
import torch
import numpy as np
import gradio as gr
import spaces # Ensure spaces is installed if needed for GPU decorator
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel, AutoConfig
import time
import re
from typing import List, Dict, Tuple, Optional
import torch.distributions as dists # Added import
import traceback # For printing exceptions

# --- START: Copied Helper functions from generation_utils.py ---
# These are needed because we are reimplementing the sampling loop locally.

def top_p_logits(logits, top_p=None):
    """ Applies top-p filtering to logits. """
    if top_p is None or top_p >= 1.0:
        return logits
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
    sorted_indices_to_remove = cumulative_probs > top_p
    # Shift the indices to the right to keep the first token above the threshold
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0

    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
    return logits

def top_k_logits(logits, top_k=None):
    """ Applies top-k filtering to logits. """
    if top_k is None or top_k <= 0:
        return logits
    top_k = min(top_k, logits.size(-1))  # Safety check
    if top_k == logits.size(-1): # Avoid unnecessary computation if k is full size
        return logits
    # Remove all tokens with a probability less than the last token of the top-k
    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
    return logits

def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False):
    """ Samples tokens based on logits and calculates confidence. """
    if temperature > 0:
        # Prevent division by zero or negative temperatures
        safe_temp = max(temperature, 1e-6)
        logits = logits / safe_temp
    if top_p is not None and 0.0 < top_p < 1.0: # Apply top_p if valid (and not disabled)
        logits = top_p_logits(logits, top_p)
    if top_k is not None and top_k > 0:    # Apply top_k if valid
        logits = top_k_logits(logits, top_k)

    # Ensure logits are not all -inf after filtering, if so, assign uniform probability.
    is_all_neg_inf = torch.all(logits <= torch.finfo(logits.dtype).min, dim=-1, keepdim=True)
    if torch.any(is_all_neg_inf):
        # print("Warning: All logits became -inf after filtering. Assigning uniform probabilities.")
        uniform_logits = torch.zeros_like(logits) # Uniform logits (zeros before softmax)
        logits = torch.where(is_all_neg_inf, uniform_logits, logits)

    probs = torch.softmax(logits, dim=-1)

    # Clamp probabilities to avoid NaNs in sampling, ensure they sum to 1
    probs = torch.clamp(probs, min=0.0) # Ensure non-negative
    prob_sum_for_norm = probs.sum(dim=-1, keepdim=True)
    # Use a tolerance check for division
    safe_prob_sum_for_norm = torch.where(prob_sum_for_norm > 1e-12, prob_sum_for_norm, torch.ones_like(prob_sum_for_norm))
    probs = probs / safe_prob_sum_for_norm # Re-normalize with safe denominator
    probs = torch.nan_to_num(probs, nan=0.0) # Handle any remaining NaNs

    if temperature > 0:
        try:
            # Ensure probs sum to 1 before sampling
            probs_sum_check = probs.sum(dim=-1)
            if not torch.all(torch.isclose(probs_sum_check, torch.ones_like(probs_sum_check))):
                 # print(f"Warning: Probs do not sum to 1 before sampling ({probs_sum_check}). Re-normalizing.")
                 probs = probs / probs.sum(dim=-1, keepdim=True) # Final normalization attempt

            x0 = dists.Categorical(probs=probs).sample()
            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
        except Exception as e: # Catch broader exceptions during sampling
            print(f"Warning: Error during Categorical sampling: {e}. Falling back to argmax.")
            confidence, x0 = probs.max(dim=-1)
    else: # Greedy decoding (temperature == 0)
        confidence, x0 = probs.max(dim=-1)

    if margin_confidence:
        sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
        # Ensure there are at least 2 probabilities to compare
        top1_probs = sorted_probs[..., 0]
        top2_probs = sorted_probs[..., 1] if sorted_probs.shape[-1] > 1 else torch.zeros_like(top1_probs) # Use 0 if only one prob
        confidence = top1_probs - top2_probs

    if neg_entropy:
        epsilon = torch.finfo(probs.dtype).eps # Use dtype's epsilon
        # Ensure probs are > 0 for log
        log_probs = torch.log(torch.clamp(probs, min=epsilon)) # Clamp before log
        confidence = torch.sum(probs * log_probs, dim=-1) # This is negative entropy

    # Ensure confidence is not NaN
    confidence = torch.nan_to_num(confidence, nan=0.0)

    return confidence, x0
# --- END: Copied Helper functions ---


# --- Model Loading and Constants ---
# Load model configuration to get special token IDs
config = AutoConfig.from_pretrained("Dream-org/Dream-v0-Instruct-7B", trust_remote_code=True)
# Use AutoModel for the base model loading, relying on trust_remote_code=True
# for the custom DreamModel class and generation mixin.
model_path = "Dream-org/Dream-v0-Instruct-7B"

# Determine device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using device: {device}")

# Load model and tokenizer
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
print("Loading model...")
# Ensure torch_dtype is set appropriately for your hardware if needed
model = AutoModel.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16 if device == 'cuda' else torch.float32, # Use bfloat16 only on CUDA
    trust_remote_code=True,
    attn_implementation="sdpa" # Explicitly request SDPA if available/desired
)
model = model.to(device).eval()
print("Model loaded.")

# Constants from Dream's config/tokenizer
MASK_TOKEN = tokenizer.mask_token
MASK_ID = tokenizer.mask_token_id # Use tokenizer's mask_token_id directly
PAD_ID = tokenizer.pad_token_id # Use tokenizer's pad_token_id
EOS_ID = tokenizer.eos_token_id # Use tokenizer's eos_token_id

if MASK_ID is None:
    print("Warning: Mask token ID not found in config/tokenizer. Trying to fetch from tokenizer...")
    mask_token_special = tokenizer.mask_token
    if mask_token_special:
        MASK_ID = tokenizer.convert_tokens_to_ids(mask_token_special)
        print(f"Found MASK_ID from tokenizer: {MASK_ID}")
    else:
        raise ValueError("Cannot determine MASK_ID. Check model's tokenizer configuration.")

SPECIAL_TOKEN_IDS = {PAD_ID, EOS_ID, MASK_ID}
try:
    IM_START_ID = tokenizer.convert_tokens_to_ids("<|im_start|>")
    IM_END_ID = tokenizer.convert_tokens_to_ids("<|im_end|>")
    if IM_START_ID is not None: SPECIAL_TOKEN_IDS.add(IM_START_ID)
    if IM_END_ID is not None: SPECIAL_TOKEN_IDS.add(IM_END_ID)
except KeyError:
    print("Warning: <|im_start|> or <|im_end|> not found in tokenizer vocab.")
    IM_START_ID = None
    IM_END_ID = None


# --- App Helper Functions ---
def parse_constraints(constraints_text: str) -> Dict[int, List[int]]:
    """ Parses constraints. """
    constraints = {}
    if not constraints_text:
        return constraints

    # Simple split on comma, assumes format 'pos:word, pos:word'
    parts = constraints_text.split(',')

    for part in parts:
        part = part.strip()
        if ':' not in part:
            continue
        pos_str, word = part.split(':', 1)
        try:
            pos = int(pos_str.strip())
            word = word.strip()
            token_ids = []
            if word: # Only encode if word is not empty
                 # Add space prefix automatically if pos > 0 and word doesn't start with space
                text_to_encode = (" " + word) if (pos > 0 and not word.startswith(" ")) else word
                token_ids = tokenizer.encode(text_to_encode, add_special_tokens=False)

            if token_ids and pos >= 0:
                constraints[pos] = token_ids
            elif not token_ids and word: # Don't warn for empty words after split
                print(f"Warning: Could not tokenize constraint word '{word}'")
        except ValueError:
            print(f"Warning: Invalid position '{pos_str}' in constraint part '{part}'")
            continue # Ignore malformed constraint parts
        except Exception as e:
            print(f"Warning: Error processing constraint '{part}': {e}")
            continue

    # print(f"Parsed constraints: {constraints}") # Debugging
    return constraints


def format_chat_history(history: List[List[Optional[str]]]) -> List[Dict[str, str]]:
    """ Formats chat history for the template. """
    messages = []
    for user_msg, assistant_msg in history:
        if user_msg is not None: # Check for None explicitly
             messages.append({"role": "user", "content": user_msg})
        # Add assistant message only if it exists (it won't for the last turn before generation)
        if assistant_msg is not None:
            messages.append({"role": "assistant", "content": assistant_msg})
    return messages

def apply_constraints_to_state(
    x: torch.Tensor,
    prompt_length: int,
    total_length: int,
    parsed_constraints: Dict[int, List[int]],
    current_step: Optional[int] = None # For logging/debugging
) -> torch.Tensor:
    """ Applies constraints directly to the state tensor `x`. """
    modified_x = x.clone() # Work on a copy
    for rel_pos, word_token_ids in parsed_constraints.items():
        abs_start_pos = prompt_length + rel_pos
        abs_end_pos = abs_start_pos + len(word_token_ids)

        # Ensure the constraint fits within the generation length
        if abs_start_pos < total_length and abs_end_pos <= total_length:
            try:
                constraint_tensor = torch.tensor(word_token_ids, dtype=torch.long, device=modified_x.device)
                # Force the constraint tokens onto the sequence
                modified_x[0, abs_start_pos:abs_end_pos] = constraint_tensor
            except IndexError:
                 print(f"Warning (Step {current_step}): Constraint at {rel_pos} ('{tokenizer.decode(word_token_ids)}') goes out of bounds.")
            except Exception as e:
                 print(f"Warning (Step {current_step}): Failed to apply constraint at {rel_pos}: {e}")
    return modified_x


# --- Core Generation Logic with Live Visualization ---

@spaces.GPU # Decorator for Hugging Face Spaces GPU usage
@torch.no_grad() # Ensure no gradients are computed during generation
def generate_dream_response(
    history: List[List[Optional[str]]], # Receives the latest state from _chat_history_store
    gen_length: int,
    steps: int,
    constraints_text: str,
    temperature: float,
    top_p: Optional[float],
    top_k: Optional[int],
    alg: str,
    alg_temp: Optional[float],
    visualization_delay: float
    ) -> List[Tuple[str, str]]:
    """ Generates text step-by-step and yields visualization states live. """

    if not history or history[-1][0] is None: # Check if last user message is None or missing
        yield history, [("Internal Error: History state invalid.", "red")], ""
        return

    # --- 1. Preparation ---
    # History already contains the latest user message and None for the bot response
    messages_for_template = format_chat_history(history)
    parsed_constraints = parse_constraints(constraints_text)

    try:
        inputs = tokenizer.apply_chat_template(
            messages_for_template,
            return_tensors="pt",
            return_dict=True,
            add_generation_prompt=True # Creates the '<|im_start|>assistant\n' prompt
        )
        input_ids = inputs.input_ids.to(device)
        # Ensure prompt_attention_mask is also on the correct device and handle missing mask
        prompt_attention_mask = inputs.attention_mask.to(device) if 'attention_mask' in inputs else torch.ones_like(input_ids)
        prompt_length = input_ids.shape[1]
    except Exception as e:
        print(f"Error applying chat template: {e}")
        # Yield current history (with None), error message, empty text
        yield history, [("Error preparing input.", "red")], ""
        return

    eps = 1e-3
    top_p_val = top_p if top_p is not None and 0.0 < top_p < 1.0 else None
    top_k_val = top_k if top_k is not None and top_k > 0 else None
    alg_temp_val = alg_temp if alg in ['maskgit_plus', 'topk_margin', 'entropy'] and alg_temp is not None and alg_temp > 0 else None

    # --- 2. Initialize Generation State ---
    total_length = prompt_length + gen_length
    initial_generation_part = torch.full((1, gen_length), MASK_ID, dtype=torch.long, device=device)
    x = torch.cat((input_ids, initial_generation_part), dim=1)

    # Prepare attention mask for SDPA (float format)
    generation_attention_mask = torch.ones((1, gen_length), dtype=prompt_attention_mask.dtype, device=device) # Match dtype
    full_attention_mask_long = torch.cat((prompt_attention_mask, generation_attention_mask), dim=1) # Shape [B, N]

    attention_mask_for_model = full_attention_mask_long.to(model.dtype) # Convert to model's float dtype
    large_neg_val = torch.finfo(model.dtype).min
    attention_mask_for_model = (1.0 - attention_mask_for_model) * large_neg_val
    attention_mask_for_model = attention_mask_for_model.unsqueeze(1).unsqueeze(2) # Shape [B, 1, 1, N]

    # Timesteps
    timesteps = torch.linspace(1, eps, steps + 1, device=device)

    # Apply initial constraints
    x = apply_constraints_to_state(x, prompt_length, total_length, parsed_constraints, current_step=-1)

    # --- 3. Visualization Setup ---
    previous_tokens_vis = None
    final_response_text = ""
    # Work on a copy of the history list received as input
    history_copy = [list(item) for item in history]

    # --- 4. Initial Yield (Masked State) ---
    initial_generated_tokens = x[0, prompt_length:].cpu()
    vis_data_initial = []
    for tok_id in initial_generated_tokens.tolist():
        display_token = MASK_TOKEN
        color = "#444444" # Dark Gray for masks
        vis_data_initial.append((display_token, color))

    previous_tokens_vis = initial_generated_tokens
    # Yield the initial history copy (with None placeholder), initial vis, empty text
    yield history_copy, vis_data_initial, ""
    time.sleep(visualization_delay)

    # --- 5. Step-by-Step Diffusion Loop ---
    try:
        start_time = time.time()
        for i in range(steps):
            mask_index = (x == MASK_ID)
            if not mask_index.any():
                 print(f"No mask tokens left at step {i}. Stopping early.")
                 break

            # --- Model Forward Pass ---
            outputs = model(
                input_ids=x,
                attention_mask=attention_mask_for_model, # Pass the [B, 1, 1, N] float mask
                position_ids=None, # Let model compute default positions
                use_cache=False,
                return_dict=True
            )
            logits = outputs.logits

            # Align logits with the token positions they predict (logits[t] predicts token[t+1])
            # Shift left, effectively aligning logits[t] with inputs[t]
            logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)

            # Select logits for masked positions
            mask_logits = logits[mask_index] # Shape [num_masked_tokens, V]
            if mask_logits.numel() == 0:
                 print(f"No masked tokens found for logit selection at step {i}. Stopping.")
                 break

            # --- Sampling / Remasking Logic ---
            t = timesteps[i]
            s = timesteps[i + 1]
            # Initialize the update tensor for masked positions with MASK_ID
            x_new_masked_part = torch.full_like(x[mask_index], MASK_ID, device=device, dtype=torch.long)

            if alg == 'origin':
                p_transfer = (1.0 - s / t) if i < steps - 1 else 1.0
                num_masked = mask_logits.shape[0]
                transfer_indices_relative = torch.rand(num_masked, device=device) < p_transfer
                logits_to_sample = mask_logits[transfer_indices_relative]

                if logits_to_sample.numel() > 0:
                    _, sampled_tokens = sample_tokens(logits_to_sample, temperature=temperature, top_p=top_p_val, top_k=top_k_val)
                    # Place sampled tokens into the correct positions within the masked part update
                    x_new_masked_part[transfer_indices_relative] = sampled_tokens

            else: # Confidence-based algorithms ('maskgit_plus', 'topk_margin', 'entropy')
                use_margin = (alg == 'topk_margin')
                use_entropy = (alg == 'entropy')
                # Sample candidates and get confidence for all masked positions
                confidence, x0_candidates = sample_tokens(
                    mask_logits,
                    temperature=temperature,
                    top_p=top_p_val,
                    top_k=top_k_val,
                    margin_confidence=use_margin,
                    neg_entropy=use_entropy
                )

                num_mask_token = mask_logits.shape[0]
                # Calculate target number of tokens to reveal in this step
                target_num_revealed_float = num_mask_token * (1.0 - s / t)
                number_transfer_tokens = int(target_num_revealed_float) if i < steps - 1 else num_mask_token

                if number_transfer_tokens > 0:
                    # Determine which tokens to reveal based on confidence
                    num_samples = min(number_transfer_tokens, num_mask_token) # Ensure k <= num_mask_token
                    if num_samples > 0:
                        transfer_indices_relative = torch.tensor([], dtype=torch.long, device=device) # Initialize empty
                        if alg_temp_val is None or alg_temp_val <= 0: # Use top-k confidence sorting
                            # Sort by confidence (higher is better, except for entropy where lower is better)
                            sort_metric = confidence if alg != 'entropy' else -confidence
                            # Ensure k is not greater than the number of elements
                            k_topk = min(num_samples, sort_metric.numel())
                            if k_topk > 0:
                                _, transfer_indices_relative = torch.topk(sort_metric, k=k_topk)

                        else: # Sample based on confidence temperature
                            # Ensure confidence has elements before processing
                            if confidence.numel() > 0:
                                conf_probs = confidence / alg_temp_val
                                # Handle potential inf/-inf before softmax, ensure non-negative probabilities
                                conf_probs = torch.nan_to_num(conf_probs, nan=0.0, posinf=1e9, neginf=-1e9)
                                # Clamp to prevent large positive values causing overflow in exp
                                conf_probs = torch.clamp(conf_probs - conf_probs.max(), min=-30) # Softmax is invariant to shift
                                conf_probs = F.softmax(conf_probs, dim=-1)
                                conf_probs = torch.clamp(conf_probs, min=0.0) # Ensure non-negative
                                conf_probs = torch.nan_to_num(conf_probs, nan=0.0) # Handle NaNs

                                # Normalize probabilities if they don't sum to 1 (within tolerance)
                                prob_sum = conf_probs.sum()
                                target_sum_tensor = torch.tensor(1.0, device=device, dtype=prob_sum.dtype)
                                if not torch.isclose(prob_sum, target_sum_tensor, atol=1e-4) and prob_sum > 0:
                                    safe_prob_sum = torch.max(prob_sum, torch.tensor(1e-12, device=device, dtype=prob_sum.dtype))
                                    conf_probs = conf_probs / safe_prob_sum

                                # Check if probabilities are valid for multinomial sampling
                                final_prob_sum_check = conf_probs.sum()
                                if conf_probs.numel() > 0 and num_samples > 0 and torch.all(conf_probs >= 0) and torch.isclose(final_prob_sum_check, target_sum_tensor, atol=1e-4):
                                    try:
                                        transfer_indices_relative = torch.multinomial(conf_probs, num_samples=num_samples, replacement=False)
                                    except RuntimeError as e:
                                        print(f"Warning step {i}: Multinomial sampling failed ('{e}'). Falling back to top-k.")
                                        # Fallback to top-k if multinomial fails
                                        sort_metric = confidence if alg != 'entropy' else -confidence
                                        k_multinomial_fallback = min(num_samples, sort_metric.numel())
                                        if k_multinomial_fallback > 0:
                                             _, transfer_indices_relative = torch.topk(sort_metric, k=k_multinomial_fallback)
                                else: # Handle cases where multinomial is not possible (e.g., bad probabilities)
                                    # print(f"Warning step {i}: Invalid probabilities for multinomial sampling (sum={final_prob_sum_check:.4f}). Falling back to top-k.")
                                    sort_metric = confidence if alg != 'entropy' else -confidence
                                    k_multinomial_fallback = min(num_samples, sort_metric.numel())
                                    if k_multinomial_fallback > 0:
                                        _, transfer_indices_relative = torch.topk(sort_metric, k=k_multinomial_fallback)

                        # Apply the transfer using the selected indices, with safety checks
                        if transfer_indices_relative.numel() > 0:
                             # Bounds check before indexing
                             max_cand_idx = x0_candidates.shape[0] - 1
                             max_mask_idx = x_new_masked_part.shape[0] - 1
                             valid_indices_mask = (transfer_indices_relative >= 0) & \
                                                  (transfer_indices_relative <= max_cand_idx) & \
                                                  (transfer_indices_relative <= max_mask_idx)
                             valid_transfer_indices = transfer_indices_relative[valid_indices_mask]

                             if valid_transfer_indices.numel() > 0:
                                  x_new_masked_part[valid_transfer_indices] = x0_candidates[valid_transfer_indices].clone()
                             # else:
                             #    if transfer_indices_relative.numel() > 0: # Only warn if there were indices initially
                             #         print(f"Warning step {i}: No valid transfer indices after bounds check.")


            # Update the global state `x` only at the masked positions
            x[mask_index] = x_new_masked_part

            # --- Apply Constraints ---
            # Constraints should be applied *after* sampling/revealing tokens for the step
            x = apply_constraints_to_state(x, prompt_length, total_length, parsed_constraints, current_step=i)

            # --- Yield Visualization ---
            current_generated_tokens = x[0, prompt_length:].cpu() # Get generated part, move to CPU
            vis_data = []
            for j in range(gen_length):
                current_tok_id = current_generated_tokens[j].item()
                # Ensure previous_tokens_vis exists and index is valid
                previous_tok_id = previous_tokens_vis[j].item() if previous_tokens_vis is not None and j < len(previous_tokens_vis) else MASK_ID

                try:
                    # Use replace='�' to handle potential bytes rendering issues in Gradio HighlightedText
                    decoded_token = tokenizer.decode([current_tok_id], skip_special_tokens=False, clean_up_tokenization_spaces=False)
                    display_token = MASK_TOKEN if current_tok_id == MASK_ID else decoded_token
                except Exception:
                    display_token = f"[ID:{current_tok_id}]" # Fallback

                color = None
                token_to_display = display_token

                if current_tok_id == MASK_ID:
                    color = "#444444" # Dark Gray for masks
                elif previous_tok_id == MASK_ID: # Token was just revealed
                    color = "#66CC66" # Light Green
                else: # Token was already revealed
                    color = "#6699CC" # Light Blue

                # Hide special tokens (PAD/EOS) if they were already revealed (LLaDA effect)
                # Ensure PAD_ID and EOS_ID are not None before checking
                should_hide = False
                if PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True
                if EOS_ID is not None and current_tok_id == EOS_ID: should_hide = True
                # Special check: If PAD and EOS are the same, only hide if it's that ID
                if PAD_ID == EOS_ID and PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True

                if should_hide and previous_tok_id == current_tok_id:
                    token_to_display = "" # Hide by making empty
                    color = None # No color for hidden

                if token_to_display: # Avoid adding empty strings if hiding
                    vis_data.append((token_to_display, color))

            # Update previous state for the next iteration's color logic
            previous_tokens_vis = current_generated_tokens

            # Decode intermediate response text using the *current* state x
            intermediate_response_tokens = x[0, prompt_length:]
            intermediate_response_text = tokenizer.decode(
                intermediate_response_tokens,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True
            ).strip()

            # Update the *copy* of the history with the intermediate text for display purposes
            if history_copy: # Ensure history_copy is not empty
                 history_copy[-1][1] = intermediate_response_text # Update the None placeholder

            # Yield the updated history copy, current vis, and intermediate text
            yield history_copy, vis_data, intermediate_response_text
            time.sleep(visualization_delay)

        end_time = time.time()
        print(f"Dream generation finished in {end_time - start_time:.2f} seconds.")

        # --- 6. Final Processing & Yield ---
        final_sequence = x[0]
        response_tokens = final_sequence[prompt_length:]
        final_response_text = tokenizer.decode(
            response_tokens,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True
        ).strip()

        # Update the final history copy *definitively*
        if history_copy:
            history_copy[-1][1] = final_response_text

        # Format the final visualization state
        final_generated_tokens = x[0, prompt_length:].cpu()
        vis_data_final = []
        for j in range(gen_length):
            current_tok_id = final_generated_tokens[j].item()
            previous_tok_id = previous_tokens_vis[j].item() if previous_tokens_vis is not None and j < len(previous_tokens_vis) else MASK_ID
            try:
                decoded_token = tokenizer.decode([current_tok_id], skip_special_tokens=False, clean_up_tokenization_spaces=False)
                display_token = MASK_TOKEN if current_tok_id == MASK_ID else decoded_token
            except Exception:
                display_token = f"[ID:{current_tok_id}]" # Fallback
            color = None
            token_to_display = display_token
            if current_tok_id == MASK_ID: color = "#444444"
            elif previous_tok_id == MASK_ID: color = "#66CC66"
            else: color = "#6699CC"

            should_hide = False
            if PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True
            if EOS_ID is not None and current_tok_id == EOS_ID: should_hide = True
            if PAD_ID == EOS_ID and PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True

            if should_hide and previous_tok_id == current_tok_id:
                 token_to_display = ""; color = None
            if token_to_display: vis_data_final.append((token_to_display, color))

        # Yield the final history, final visualization, and final text
        yield history_copy, vis_data_final, final_response_text
        print("Visualization streaming complete.")

    except Exception as e:
        print(f"Error during generation or processing loop: {e}")
        traceback.print_exc()
        # Yield the history as it was before the error, error vis, empty text
        yield history_copy, [("Error during generation.", "red")], ""
        return


# --- Gradio UI ---
css = '''
.category-legend{display:none}
button{min-height: 60px}
'''
def create_chatbot_demo():
    with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Dream 7B - Diffusion Language Model Demo")
        gr.Markdown(
            "[[Model Card](https://huggingface.co/Dream-org/Dream-v0-Instruct-7B)] "
            "[[Blog](https://hkunlp.github.io/blog/2025/dream/)]" # Note: Link might be hypothetical
        )

        # STATE MANAGEMENT
        _chat_history_store = gr.State([]) # Hidden state to store actual history list

        # UI COMPONENTS
        with gr.Row():
            with gr.Column(scale=3):
                chatbot_ui = gr.Chatbot(
                    label="Conversation",
                    height=500,
                    show_copy_button=True,
                    bubble_full_width=False,
                )
                with gr.Group():
                    with gr.Row():
                        user_input = gr.Textbox(
                            label="Your Message",
                            placeholder="Type your message here...",
                            scale=7,
                            autofocus=True,
                            show_label=False,
                            container=False
                        )
                        send_btn = gr.Button("Send", scale=1, variant="primary")
                constraints_input = gr.Textbox(
                    label="Word Constraints (Optional)",
                    info="Place words at specific positions (0-indexed from start of generation). Format: 'pos:word, pos:word,...'. Example: '0:Once, 5:upon, 10:time'",
                    placeholder="0:Hello, 10:world",
                    value=""
                )
            with gr.Column(scale=2):
                output_vis = gr.HighlightedText(
                    label="Denoising Process Visualization",
                    combine_adjacent=False,
                    show_legend=True,
                    interactive=False,
                )
                response_text_display = gr.Textbox(
                    label="Generated Response",
                    interactive=False,
                    lines=5
                )

        # Advanced generation settings
        with gr.Accordion("Generation Settings", open=False):
             with gr.Row():
                gen_length = gr.Slider(minimum=16, maximum=512, value=128, step=8, label="Max New Tokens")
                steps = gr.Slider(minimum=8, maximum=512, value=128, step=8, label="Diffusion Steps")
             with gr.Row():
                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.05, label="Temperature (0 = greedy)")
                alg_temp = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.05, label="Remasking Temp (Confidence Algs)")
             with gr.Row():
                # Adjusted label for clarity on disabling top_p
                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="Top-P (>0 & <1 to enable)")
                top_k = gr.Slider(minimum=0, maximum=200, value=0, step=5, label="Top-K (>0 to enable)")
             with gr.Row():
                 remasking_strategy = gr.Radio(choices=['origin', 'maskgit_plus', 'topk_margin', 'entropy'], value='entropy', label="Remasking Strategy (Algorithm)")
             with gr.Row():
                visualization_delay = gr.Slider(minimum=0.0, maximum=0.5, value=0.03, step=0.01, label="Visualization Delay (seconds)")

        # Clear button
        clear_btn = gr.Button("Clear Conversation")

        # --- Event Handlers ---

        def add_user_message_to_history(message: str, history_store: List[List[Optional[str]]]):
            """Adds user message TO STATE, clears input, prepares for bot response."""
            if not message.strip():
                gr.Warning("Please enter a message.")
                # Return unchanged state, but clear inputs/outputs for next step
                # Outputs: _chat_history_store, user_input, output_vis, response_text_display
                return history_store, message, [], "" # Return original message to keep it in input if invalid

            # Add user message with placeholder for bot response TO THE STATE
            history_store.append([message.strip(), None]) # Ensure message is stripped
            # Return updated history store, clear input box, clear vis, clear response text
            # Outputs: _chat_history_store, user_input, output_vis, response_text_display
            return history_store, "", [], "" # Clear user_input only on success

        def clear_conversation():
            """Clears the chat history state and UI elements."""
            # Outputs: _chat_history_store, chatbot_ui, user_input, output_vis, response_text_display
            return [], [], "", [], "" # Clear everything


        # --- Connect UI elements ---

        # Inputs for the generation function
        generation_inputs = [
            _chat_history_store, gen_length, steps, constraints_input,
            temperature, top_p, top_k, remasking_strategy, alg_temp,
            visualization_delay
        ]
        # Outputs for the generation function (yields history, vis_data, text)
        generation_outputs = [chatbot_ui, output_vis, response_text_display]

        # Outputs for add_user_message_to_history
        add_message_outputs = [
            _chat_history_store, # Update state
            user_input,          # Clear input (or return original if invalid)
            output_vis,          # Clear visualization
            response_text_display # Clear response text
        ]

        # Handle Textbox Submission (Enter key)
        submit_listener = user_input.submit(
            fn=add_user_message_to_history,
            inputs=[user_input, _chat_history_store],
            outputs=add_message_outputs, # Step 1: Update state, clear inputs/vis/response
            queue=True # Ensure intermediate steps are processed
        ).then(
            fn=generate_dream_response,
            inputs=generation_inputs, # Takes the updated state
            outputs=generation_outputs, # Step 2: Generate response and stream history/vis/text to UI
            show_progress="hidden", # Hide default progress as we have live vis
            queue=True # Ensure generation runs in the queue
        )

        # Handle Send Button Click
        click_listener = send_btn.click(
            fn=add_user_message_to_history,
            inputs=[user_input, _chat_history_store],
            outputs=add_message_outputs, # Step 1: Update state, clear inputs/vis/response
            queue=True # Ensure intermediate steps are processed
        ).then(
            fn=generate_dream_response,
            inputs=generation_inputs, # Takes the updated state
            outputs=generation_outputs, # Step 2: Generate response and stream history/vis/text to UI
            show_progress="hidden", # Hide default progress as we have live vis
            queue=True # Ensure generation runs in the queue
        )

        # Clear Button Action
        clear_btn.click(
            clear_conversation,
            inputs=[],
            outputs=[_chat_history_store, chatbot_ui, user_input, output_vis, response_text_display],
            queue=False # Clearing can be immediate
        )

    return demo

# --- Launch ---
if __name__ == "__main__":
    demo = create_chatbot_demo()
    # Use queue for handling multiple users and streaming
    demo.queue().launch(debug=True, share=False) # Set share=True for public link