import gradio as gr import torch import numpy as np import torch.nn.functional as F from transformers import AutoTokenizer from torchvision import transforms from models import MAGVITv2, get_mask_schedule, MMadaModelLM from training.prompting_utils import UniversalPrompting from PIL import Image import spaces # --- 辅助函数 (未修改) --- def image_transform(image, resolution=256, normalize=True): image = transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BICUBIC)(image) image = transforms.CenterCrop((resolution, resolution))(image) image = transforms.ToTensor()(image) if normalize: image = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)(image) return image def add_gumbel_noise(logits, temperature): if abs(temperature) < 1e-9: return logits logits = logits.to(torch.float64) noise = torch.rand_like(logits, dtype=torch.float64) standard_gumbel_noise = -torch.log(-torch.log(noise + 1e-20) + 1e-20) return logits + temperature * standard_gumbel_noise def get_num_transfer_tokens(mask_index, steps): mask_num = mask_index.sum(dim=1, keepdim=True) steps = max(1, int(steps)) base = mask_num // steps remainder = mask_num % steps num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.long) + base for i in range(mask_num.size(0)): if remainder[i] > 0 : num_transfer_tokens[i, :remainder[i].item()] += 1 return num_transfer_tokens # --- 全局变量和模型配置 --- DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' # 固定使用 MMaDA-8B-MixCoT 模型 DEFAULT_MODEL_PATH = "Gen-Verse/MMaDA-8B-MixCoT" MASK_ID = None MODEL = None TOKENIZER = None uni_prompting = None VQ_MODEL = None # --- 核心模型加载函数 (已简化) --- @spaces.GPU def load_model_and_tokenizer(): """ 加载固定的 MMaDA-8B-MixCoT 模型和分词器。 """ global MODEL, TOKENIZER, MASK_ID, DEVICE, uni_prompting # 如果模型已经加载，则直接返回 if MODEL is not None: return f"Model 'MMaDA-8B-MixCoT' is already loaded. MASK_ID: {MASK_ID}" status_msg_parts = [f"Loading 'MMaDA-8B-MixCoT'..."] try: TOKENIZER = AutoTokenizer.from_pretrained(DEFAULT_MODEL_PATH, trust_remote_code=True) status_msg_parts.append(f"Tokenizer for 'MMaDA-8B-MixCoT' loaded.") MODEL = MMadaModelLM.from_pretrained(DEFAULT_MODEL_PATH, trust_remote_code=True, torch_dtype=torch.bfloat16).eval() status_msg_parts.append(f"Model 'MMaDA-8B-MixCoT' loaded to {DEVICE}.") uni_prompting = UniversalPrompting(TOKENIZER, max_text_len=512, special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"),ignore_id=-100, cond_dropout_prob=0.1, use_reserved_token=True) MASK_ID = 126336 status_msg_parts.append(f"Using default MASK_ID: {MASK_ID}.") if TOKENIZER.pad_token_id is None: if TOKENIZER.eos_token_id is not None: TOKENIZER.pad_token_id = TOKENIZER.eos_token_id TOKENIZER.pad_token = TOKENIZER.eos_token status_msg_parts.append(f"Set pad_token_id to eos_token_id ({TOKENIZER.eos_token_id}).") TOKENIZER.chat_template = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n' }}" return " ".join(status_msg_parts) except Exception as e: MODEL, TOKENIZER, MASK_ID = None, None, None return f"Error loading model 'MMaDA-8B-MixCoT': {str(e)}" # --- 可视化和生成函数 (generate_viz_wrapper* 系列，已修复全局变量问题) --- def get_highlighted_text_tuples(current_x_ids_batch, prompt_input_ids, prompt_len, tk, current_mask_id, raw_prompt_attention_mask): if current_x_ids_batch is None or current_x_ids_batch.ndim == 0 or current_x_ids_batch.shape[0] == 0: return [("Error in sequence data for visualization.", "ERROR")] current_x_ids_batch = current_x_ids_batch[:, prompt_len:] seq_ids = current_x_ids_batch[0].tolist() intermediate_tuples = [] for j, token_id_int in enumerate(seq_ids): try: token_str = tk.decode([token_id_int], skip_special_tokens=True, clean_up_tokenization_spaces=False) except Exception: token_str = f"[ID:{token_id_int}]" label = "ERROR" if token_id_int == current_mask_id: token_str = "[MASK]" label = "MASK" else: label = "GEN" intermediate_tuples.append((token_str, label, token_id_int)) return intermediate_tuples @torch.no_grad() @spaces.GPU def generate_viz_wrapper_t2i(prompt_text, steps, guidance_scale, mask_schedule="cosine"): global MODEL, TOKENIZER, MASK_ID, DEVICE, uni_prompting, VQ_MODEL if MODEL is None or TOKENIZER is None or MASK_ID is None: yield Image.new("RGB", (512, 512), (255, 255, 255)), "Error: Model not loaded. Please load the model first." return if DEVICE == 'cuda': MODEL.to(DEVICE) VQ_MODEL.to(DEVICE) try: # ... (函数实现和之前一样) steps = int(steps) guidance_scale = float(guidance_scale) image_tokens = torch.ones((1, 1024), dtype=torch.long, device=DEVICE) * MASK_ID prompt_text = [prompt_text] input_ids, attention_mask = uni_prompting((prompt_text, image_tokens), 't2i_gen') if guidance_scale > 0: uncond_input_ids, uncond_attention_mask = uni_prompting(([''], image_tokens), 't2i_gen') else: uncond_input_ids, uncond_attention_mask = None, None mask_schedule = get_mask_schedule(mask_schedule) blank_image = Image.new("RGB", (512, 512), (255, 255, 255)) yield blank_image, "Starting generation..." for image_step, status_msg_step in MODEL.t2i_generate_decoding_stepwise( input_ids=input_ids, uncond_input_ids=uncond_input_ids, attention_mask=attention_mask, uncond_attention_mask=uncond_attention_mask, temperature=1.0, timesteps=steps, guidance_scale=guidance_scale, noise_schedule=mask_schedule, noise_type="mask", seq_len=1024, vq_model=VQ_MODEL, uni_prompting=uni_prompting): yield image_step, status_msg_step finally: if DEVICE == 'cuda': MODEL.to('cpu') VQ_MODEL.to('cpu') torch.cuda.empty_cache() @torch.no_grad() @spaces.GPU def generate_viz_wrapper_lm(prompt_text, steps, gen_length, block_length, temperature, cfg_scale, remasking_strategy, thinking_mode_lm=False): global MODEL, TOKENIZER, MASK_ID, DEVICE if MODEL is None or TOKENIZER is None or MASK_ID is None: yield [("Error: Model not loaded. Please load the model first.", "ERROR")], "Model not loaded." return if DEVICE == 'cuda': MODEL.to(DEVICE) try: # ... (函数实现和之前一样) steps, gen_length, block_length = int(steps), int(gen_length), int(block_length) if thinking_mode_lm: prompt_text = "You should first think about the reasoning process in the mind and then provide the user with the answer. The reasoning process is enclosed within tags, i.e. reasoning process here answer here\n" + prompt_text m = [{"role": "user", "content": prompt_text}] processed_prompt_text = TOKENIZER.apply_chat_template(m, add_generation_prompt=True, tokenize=False) input_ids = TOKENIZER(text=processed_prompt_text, return_tensors="pt", padding="longest", padding_side="left", truncation=True, max_length=4096)['input_ids'].to(DEVICE) raw_prompt_attention_mask = torch.ones_like(input_ids) # Dummy mask, adjust if needed batch_size, prompt_len = input_ids.shape[0], input_ids.shape[1] x = torch.full((batch_size, prompt_len + gen_length), MASK_ID, dtype=torch.long, device=DEVICE) x[:, :prompt_len] = input_ids.clone() yield get_highlighted_text_tuples(x, input_ids, prompt_len, TOKENIZER, MASK_ID, raw_prompt_attention_mask), "Starting generation..." # ... (rest of the logic is the same) num_blocks = gen_length // block_length steps_per_block = steps // num_blocks for num_block_iter in range(num_blocks): current_block_start_idx_in_x = prompt_len + num_block_iter * block_length current_block_end_idx_in_x = prompt_len + (num_block_iter + 1) * block_length block_masks_bool_current = torch.zeros_like(x, dtype=torch.bool) block_masks_bool_current[:, current_block_start_idx_in_x:current_block_end_idx_in_x] = (x[:, current_block_start_idx_in_x:current_block_end_idx_in_x] == MASK_ID) num_transfer_tokens_for_this_block = get_num_transfer_tokens(block_masks_bool_current[:, current_block_start_idx_in_x:current_block_end_idx_in_x], steps_per_block) for i_step_in_block in range(steps_per_block): mask_index_global = (x == MASK_ID) model_output = MODEL(x) logits = model_output.logits logits_with_noise = add_gumbel_noise(logits, temperature=temperature) x0_predicted_tokens = torch.argmax(logits_with_noise, dim=-1) probs = F.softmax(logits.to(torch.float64), dim=-1) x0_probs = torch.gather(probs, dim=-1, index=x0_predicted_tokens.unsqueeze(-1)).squeeze(-1) confidence_for_selection = torch.where(mask_index_global & block_masks_bool_current, x0_probs, -torch.inf) x0_final_candidates = torch.where(mask_index_global, x0_predicted_tokens, x) transfer_indices_bool = torch.zeros_like(x, dtype=torch.bool) num_to_transfer_this_step_batch = num_transfer_tokens_for_this_block[:, i_step_in_block] for j_batch_idx in range(batch_size): k_val = min(num_to_transfer_this_step_batch[j_batch_idx].item(), candidate_positions_for_unmasking[j_batch_idx].sum().item()) if k_val > 0: _, topk_indices_in_x = torch.topk(confidence_for_selection[j_batch_idx], k=k_val) transfer_indices_bool[j_batch_idx, topk_indices_in_x] = True x[transfer_indices_bool] = x0_final_candidates[transfer_indices_bool] status_msg = f"Block {num_block_iter+1}/{num_blocks}, Step {i_step_in_block+1}/{steps_per_block}" yield get_highlighted_text_tuples(x, input_ids, prompt_len, TOKENIZER, MASK_ID, raw_prompt_attention_mask), status_msg final_text_output = TOKENIZER.batch_decode(x[:, prompt_len:], skip_special_tokens=True) yield get_highlighted_text_tuples(x, input_ids, prompt_len, TOKENIZER, MASK_ID, raw_prompt_attention_mask), final_text_output[0] finally: if DEVICE == 'cuda': MODEL.to('cpu') torch.cuda.empty_cache() @torch.no_grad() @spaces.GPU def generate_viz_wrapper(uploaded_image_pil, prompt_text, steps, gen_length, block_length, temperature, cfg_scale, remasking_strategy, thinking_mode_mmu=False): global MODEL, TOKENIZER, MASK_ID, DEVICE, VQ_MODEL if MODEL is None or TOKENIZER is None or MASK_ID is None: yield [("Error: Model not loaded. Please load the model first.", "ERROR")], "Model not loaded." return if DEVICE == 'cuda': MODEL.to(DEVICE) VQ_MODEL.to(DEVICE) try: # ... (函数实现和之前一样) steps, gen_length, block_length = int(steps), int(gen_length), int(block_length) if thinking_mode_mmu: prompt_text = "You should first think about the reasoning process in the mind and then provide the user with the answer. The reasoning process is enclosed within tags, i.e. reasoning process here answer here\n" + prompt_text m = [{"role": "user", "content": prompt_text}] processed_prompt_text = TOKENIZER.apply_chat_template(m, add_generation_prompt=True, tokenize=False) image_vq_ids_tensor = None if uploaded_image_pil is not None: image = image_transform(uploaded_image_pil, resolution=512).to(DEVICE).unsqueeze(0) image_vq_ids_tensor = VQ_MODEL.get_code(image) + 126349 input_ids = TOKENIZER(text=processed_prompt_text, return_tensors="pt", padding="longest", padding_side="left", truncation=True, max_length=4096)['input_ids'].to(DEVICE) raw_prompt_attention_mask = torch.ones_like(input_ids) # Dummy mask if image_vq_ids_tensor is not None: input_ids = torch.cat([(torch.ones(1, 1) * 126089).to(DEVICE), (torch.ones(1, 1) * 126084).to(DEVICE), image_vq_ids_tensor, (torch.ones(1, 1) * 126085).to(DEVICE), input_ids], dim=1).long() batch_size, prompt_len = input_ids.shape[0], input_ids.shape[1] x = torch.full((batch_size, prompt_len + gen_length), MASK_ID, dtype=torch.long, device=DEVICE) x[:, :prompt_len] = input_ids.clone() yield get_highlighted_text_tuples(x, input_ids, prompt_len, TOKENIZER, MASK_ID, raw_prompt_attention_mask), "Starting generation..." # ... (rest of the logic is the same) num_blocks = gen_length // block_length steps_per_block = steps // num_blocks for num_block_iter in range(num_blocks): current_block_start_idx_in_x = prompt_len + num_block_iter * block_length current_block_end_idx_in_x = prompt_len + (num_block_iter + 1) * block_length block_masks_bool_current = torch.zeros_like(x, dtype=torch.bool) block_masks_bool_current[:, current_block_start_idx_in_x:current_block_end_idx_in_x] = (x[:, current_block_start_idx_in_x:current_block_end_idx_in_x] == MASK_ID) num_transfer_tokens_for_this_block = get_num_transfer_tokens(block_masks_bool_current[:, current_block_start_idx_in_x:current_block_end_idx_in_x], steps_per_block) for i_step_in_block in range(steps_per_block): mask_index_global = (x == MASK_ID) model_output = MODEL(x) logits = model_output.logits logits_with_noise = add_gumbel_noise(logits, temperature=temperature) x0_predicted_tokens = torch.argmax(logits_with_noise, dim=-1) probs = F.softmax(logits.to(torch.float64), dim=-1) x0_probs = torch.gather(probs, dim=-1, index=x0_predicted_tokens.unsqueeze(-1)).squeeze(-1) confidence_for_selection = torch.where(mask_index_global & block_masks_bool_current, x0_probs, -torch.inf) x0_final_candidates = torch.where(mask_index_global, x0_predicted_tokens, x) transfer_indices_bool = torch.zeros_like(x, dtype=torch.bool) num_to_transfer_this_step_batch = num_transfer_tokens_for_this_block[:, i_step_in_block] for j_batch_idx in range(batch_size): k_val = min(num_to_transfer_this_step_batch[j_batch_idx].item(), (mask_index_global & block_masks_bool_current)[j_batch_idx].sum().item()) if k_val > 0: _, topk_indices_in_x = torch.topk(confidence_for_selection[j_batch_idx], k=k_val) transfer_indices_bool[j_batch_idx, topk_indices_in_x] = True x[transfer_indices_bool] = x0_final_candidates[transfer_indices_bool] status_msg = f"Block {num_block_iter+1}/{num_blocks}, Step {i_step_in_block+1}/{steps_per_block}" yield get_highlighted_text_tuples(x, input_ids, prompt_len, TOKENIZER, MASK_ID, raw_prompt_attention_mask), status_msg final_text_output = TOKENIZER.batch_decode(x[:, prompt_len:], skip_special_tokens=True) yield get_highlighted_text_tuples(x, input_ids, prompt_len, TOKENIZER, MASK_ID, raw_prompt_attention_mask), final_text_output[0] finally: if DEVICE == 'cuda': MODEL.to('cpu') VQ_MODEL.to('cpu') torch.cuda.empty_cache() # --- UI定义 --- css_styles = """ .gradio-container{font-family:'IBM Plex Sans',sans-serif;margin:auto;} .gr-input {background:#f9f9f9 !important;border:1px solid #e0e0e0 !important;} .gr-output{background:#f0f0f0 !important;border:1px solid #d0d0d0 !important;} .highlighted-text span{padding:2px 4px;border-radius:4px;margin:1px 2px;display:inline-block;line-height:1.6;} footer{display:none !important} #live-update-scrollable-box {max-height: 800px; overflow-y: auto !important; display: block;} #think_btn {background-color: #f3f4f6 !important; border: 1px solid #d0d0d0 !important; color: #111827 !important; font-size: 16px !important; font-weight: bold !important;} #think_btn:hover {background-color: #e0e0e0 !important; border: 1px solid #c0c0c0 !important; color: #222 !important;} #think_btn:active {background-color: #2563eb !important; border: 1px solid #b0b0b0 !important; color: white !important;} .model-badge {padding: 5px 10px; border-radius: 15px; font-weight: bold; margin: 0 5px; display: inline-block;} .active-model {background-color: #E879F9; color: white;} .soon-model {background-color: #E5E7EB; color: #6B7280; cursor: not-allowed;} """ def toggle_thinking_mode(current_thinking_mode): new_state = not current_thinking_mode new_label = "Thinking Mode ✅" if new_state else "Thinking Mode ❌" return new_state, gr.update(value=new_label) color_map_config = {"MASK": "lightgrey", "GEN": "#DCABFA"} theme = gr.themes.Ocean(primary_hue="fuchsia") with gr.Blocks(css=css_styles, theme=theme) as demo: thinking_mode_lm = gr.State(True) # MixCoT模型默认开启 thinking_mode_mmu = gr.State(True) # MixCoT模型默认开启 # --- 标题和模型信息 (已修改) --- gr.HTML("""

MMaDA is a new class of multimodal diffusion foundation models, enabling state-of-the-art performance in reasoning, multimodal understanding, and text-to-image generation.

📄 Paper   |   💻 Code

""") with gr.Row(): with gr.Column(scale=1): gr.HTML("""
MMaDA-8B-MixCoT MMaDA-8B-Max (coming soon)
""") with gr.Column(scale=2): model_load_status_box = gr.Textbox( label="Model Load Status", interactive=False, lines=3, max_lines=5 ) # --- Part 1. 文本生成 --- gr.Markdown("## Part 1. Text Generation") with gr.Row(): with gr.Column(scale=2): prompt_input_box_lm = gr.Textbox(label="Enter your prompt:", lines=3, value="A rectangular prism has a length of 5 units, a width of 4 units, and a height of 3 units. What is the volume of the prism?") think_button_lm = gr.Button("Thinking Mode ✅", elem_id="think_btn") with gr.Accordion("Generation Parameters", open=True): # ... 参数滑块 (未修改) with gr.Row(): gen_length_slider_lm = gr.Slider(minimum=8, maximum=1024, value=512, step=64, label="Generation Length") steps_slider_lm = gr.Slider(minimum=1, maximum=512, value=256, step=32, label="Total Sampling Steps") with gr.Row(): block_length_slider_lm = gr.Slider(minimum=8, maximum=1024, value=128, step=32, label="Block Length") remasking_dropdown_lm = gr.Dropdown(choices=['low_confidence', 'random'], value='low_confidence', label="Remasking Strategy") with gr.Row(): cfg_scale_slider_lm = gr.Slider(minimum=0.0, maximum=2.0, value=0.0, step=0.1, label="CFG Scale") temperature_slider_lm = gr.Slider(minimum=0.0, maximum=2.0, value=1, step=0.05, label="Temperature") with gr.Row(): run_button_ui_lm = gr.Button("Generate Sequence", variant="primary", scale=3) clear_button_ui_lm = gr.Button("Clear Outputs", scale=1) with gr.Column(scale=3): output_visualization_box_lm = gr.HighlightedText(label="Live Generation Process", show_legend=True, color_map=color_map_config, combine_adjacent=False, interactive=False, elem_id="live-update-scrollable-box") output_final_text_box_lm = gr.Textbox(label="Final Output", lines=8, interactive=False, show_copy_button=True) # 仅保留 MixCoT 的示例 (已修改) gr.Examples( examples=[ ["A rectangular prism has a length of 5 units, a width of 4 units, and a height of 3 units. What is the volume of the prism?", 256, 512, 128, 1, 0, "low_confidence"], ["Lily can run 12 kilometers per hour for 4 hours. After that, she can run 6 kilometers per hour. How many kilometers can she run in 8 hours?", 256, 512, 64, 1, 0, "low_confidence"] ], inputs=[prompt_input_box_lm, steps_slider_lm, gen_length_slider_lm, block_length_slider_lm, temperature_slider_lm, cfg_scale_slider_lm, remasking_dropdown_lm], outputs=[output_visualization_box_lm, output_final_text_box_lm], fn=generate_viz_wrapper_lm, cache_examples=False ) # --- Part 2 & 3 和事件处理器 (结构类似，已做简化) --- gr.Markdown("---") gr.Markdown("## Part 2. Multimodal Understanding") with gr.Row(): # ... (Part 2 UI 结构未变) with gr.Column(scale=2): prompt_input_box_mmu = gr.Textbox(label="Enter your prompt:", lines=3, value="") think_button_mmu = gr.Button("Thinking Mode ✅", elem_id="think_btn") with gr.Accordion("Generation Parameters", open=True): with gr.Row(): gen_length_slider_mmu = gr.Slider(minimum=64, maximum=1024, value=512, step=64, label="Generation Length") steps_slider_mmu = gr.Slider(minimum=1, maximum=512, value=256, step=32, label="Total Sampling Steps") with gr.Row(): block_length_slider_mmu = gr.Slider(minimum=32, maximum=1024, value=64, step=32, label="Block Length") remasking_dropdown_mmu = gr.Dropdown(choices=['low_confidence', 'random'], value='low_confidence', label="Remasking Strategy") with gr.Row(): cfg_scale_slider_mmu = gr.Slider(minimum=0.0, maximum=2.0, value=0.0, step=0.1, label="CFG Scale") temperature_slider_mmu = gr.Slider(minimum=0.0, maximum=2.0, value=1, step=0.05, label="Temperature") with gr.Row(): image_upload_box = gr.Image(type="pil", label="Upload Image") with gr.Row(): run_button_ui_mmu = gr.Button("Generate Description", variant="primary", scale=3) clear_button_ui_mmu = gr.Button("Clear Outputs", scale=1) with gr.Column(scale=3): output_visualization_box_mmu = gr.HighlightedText(label="Token Sequence (Live Update)", show_legend=True, color_map=color_map_config, combine_adjacent=False, interactive=False, elem_id="live-update-scrollable-box") output_final_text_box_mmu = gr.Textbox(label="Final Output", lines=8, interactive=False, show_copy_button=True) # 仅保留 MixCoT 的 MMU 示例 gr.Examples( examples=[ ["figs/geo.png", "In the given figure, a square ABCD is inscribed in a circle with center O. Point P is located on side CD. What is the value of angle APB?", 256, 512, 64, 1, 0, "low_confidence"], ["figs/bus.jpg", "What are the colors of the bus?", 256, 512, 64, 1, 0, "low_confidence"] ], inputs=[image_upload_box, prompt_input_box_mmu, steps_slider_mmu, gen_length_slider_mmu, block_length_slider_mmu, temperature_slider_mmu, cfg_scale_slider_mmu, remasking_dropdown_mmu], outputs=[output_visualization_box_mmu, output_final_text_box_mmu], fn=generate_viz_wrapper, cache_examples=False ) gr.Markdown("---") gr.Markdown("## Part 3. Text-to-Image Generation") # ... (Part 3 UI 和示例未变) with gr.Row(): with gr.Column(scale=2): prompt_input_box_t2i = gr.Textbox(label="Enter your prompt:", lines=3, value="A sea turtle swimming near a coral reef in the ocean, with a clear blue sky and water in the background.") with gr.Accordion("Generation Parameters", open=True): with gr.Row(): steps_slider_t2i = gr.Slider(minimum=5, maximum=100, value=15, step=5, label="Total Sampling Steps") guidance_scale_slider_t2i = gr.Slider(minimum=0.0, maximum=7.0, value=3.5, step=0.5, label="Guidance Scale") with gr.Row(): scheduler_radio_t2i = gr.Radio(choices=["cosine", "sigmoid", "linear"], value="cosine", label="Scheduler") with gr.Row(): run_button_ui_t2i = gr.Button("Generate Image", variant="primary", scale=3) clear_button_ui_t2i = gr.Button("Clear Outputs", scale=1) with gr.Column(scale=3): output_image_t2i = gr.Image(label="Generated Image", interactive=False, type="pil") output_status_t2i = gr.Textbox(label="Generation Status", interactive=False) gr.Examples( examples=[ ["A sea turtle swimming near a coral reef in the ocean, with a clear blue sky and water in the background.", 15, 3.5, "cosine"], ["A beautiful sunset over a calm ocean, with a few clouds in the sky.", 15, 3.5, "cosine"] ], inputs=[prompt_input_box_t2i, steps_slider_t2i, guidance_scale_slider_t2i, scheduler_radio_t2i], outputs=[output_image_t2i, output_status_t2i], fn=generate_viz_wrapper_t2i, cache_examples=False ) # --- 应用启动和事件处理 (已简化) --- def initialize_app_state(): global VQ_MODEL print("Loading VQ_MODEL for the first time...") VQ_MODEL = MAGVITv2().from_pretrained("showlab/magvitv2") print("VQ_MODEL loaded to CPU.") status = load_model_and_tokenizer() # MixCoT模型默认开启Thinking Mode return status, True, gr.update(value="Thinking Mode ✅"), True, gr.update(value="Thinking Mode ✅") demo.load( fn=initialize_app_state, inputs=None, outputs=[ model_load_status_box, thinking_mode_lm, think_button_lm, thinking_mode_mmu, think_button_mmu ], queue=True ) # 清除按钮事件 clear_button_ui_lm.click(fn=lambda: (None, None), inputs=None, outputs=[output_visualization_box_lm, output_final_text_box_lm], queue=False) clear_button_ui_mmu.click(fn=lambda: (None, None, None), inputs=None, outputs=[image_upload_box, output_visualization_box_mmu, output_final_text_box_mmu], queue=False) clear_button_ui_t2i.click(fn=lambda: (None, ""), inputs=None, outputs=[output_image_t2i, output_status_t2i], queue=False) # Thinking Mode 切换事件 think_button_lm.click(fn=toggle_thinking_mode, inputs=[thinking_mode_lm], outputs=[thinking_mode_lm, think_button_lm]) think_button_mmu.click(fn=toggle_thinking_mode, inputs=[thinking_mode_mmu], outputs=[thinking_mode_mmu, think_button_mmu]) # 生成按钮事件 run_button_ui_lm.click(fn=generate_viz_wrapper_lm, inputs=[prompt_input_box_lm, steps_slider_lm, gen_length_slider_lm, block_length_slider_lm, temperature_slider_lm, cfg_scale_slider_lm, remasking_dropdown_lm, thinking_mode_lm], outputs=[output_visualization_box_lm, output_final_text_box_lm]) run_button_ui_mmu.click(fn=generate_viz_wrapper, inputs=[image_upload_box, prompt_input_box_mmu, steps_slider_mmu, gen_length_slider_mmu, block_length_slider_mmu, temperature_slider_mmu, cfg_scale_slider_mmu, remasking_dropdown_mmu, thinking_mode_mmu], outputs=[output_visualization_box_mmu, output_final_text_box_mmu]) run_button_ui_t2i.click(fn=generate_viz_wrapper_t2i, inputs=[prompt_input_box_t2i, steps_slider_t2i, guidance_scale_slider_t2i, scheduler_radio_t2i], outputs=[output_image_t2i, output_status_t2i]) if __name__ == "__main__": print(f"Starting Gradio App. Attempting to use device: {DEVICE}") demo.launch(allowed_paths=["title.png", "figs"])