debug

app.py

@@ -1,30 +1,263 @@
+import spaces
 import gradio as gr
+import numpy as np
+import os
+import torch
+import random
+import subprocess
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
+from accelerate import infer_auto_device_map, load_checkpoint_and_dispatch, init_empty_weights
 from PIL import Image
 
+from data.data_utils import add_special_tokens, pil_img2rgb
+from data.transforms import ImageTransform
+from inferencer import InterleaveInferencer
+from modeling.autoencoder import load_ae
+from modeling.bagel import (
+    BagelConfig, Bagel, Qwen2Config, Qwen2ForCausalLM,
+    SiglipVisionConfig, SiglipVisionModel
+)
+from modeling.qwen2 import Qwen2Tokenizer
+
+from huggingface_hub import snapshot_download
+
+save_dir = "./model_weights"
+repo_id = "ByteDance-Seed/BAGEL-7B-MoT"
+cache_dir = save_dir + "/cache"
+
+snapshot_download(
+  cache_dir=cache_dir,
+  local_dir=save_dir,
+  repo_id=repo_id,
+  local_dir_use_symlinks=False,
+  resume_download=True,
+  allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
+)
+
+# Model Initialization
+model_path = save_dir
+
+llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
+llm_config.qk_norm = True
+llm_config.tie_word_embeddings = False
+llm_config.layer_module = "Qwen2MoTDecoderLayer"
+
+vit_config = SiglipVisionConfig.from_json_file(os.path.join(model_path, "vit_config.json"))
+vit_config.rope = False
+vit_config.num_hidden_layers -= 1
+
+vae_model, vae_config = load_ae(local_path=os.path.join(model_path, "ae.safetensors"))
+
+config = BagelConfig(
+    visual_gen=True,
+    visual_und=True,
+    llm_config=llm_config, 
+    vit_config=vit_config,
+    vae_config=vae_config,
+    vit_max_num_patch_per_side=70,
+    connector_act='gelu_pytorch_tanh',
+    latent_patch_size=2,
+    max_latent_size=64,
+)
+
+with init_empty_weights():
+    language_model = Qwen2ForCausalLM(llm_config)
+    vit_model      = SiglipVisionModel(vit_config)
+    model          = Bagel(language_model, vit_model, config)
+    model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(vit_config, meta=True)
+
+tokenizer = Qwen2Tokenizer.from_pretrained(model_path)
+tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)
+
+vae_transform = ImageTransform(1024, 512, 16)
+vit_transform = ImageTransform(980, 224, 14)
+
+# Model Loading and Multi GPU Infernece Preparing
+device_map = infer_auto_device_map(
+    model,
+    max_memory={i: "80GiB" for i in range(torch.cuda.device_count())},
+    no_split_module_classes=["Bagel", "Qwen2MoTDecoderLayer"],
+)
+
+same_device_modules = [
+    'language_model.model.embed_tokens',
+    'time_embedder',
+    'latent_pos_embed',
+    'vae2llm',
+    'llm2vae',
+    'connector',
+    'vit_pos_embed'
+]
+
+if torch.cuda.device_count() == 1:
+    first_device = device_map.get(same_device_modules[0], "cuda:0")
+    for k in same_device_modules:
+        if k in device_map:
+            device_map[k] = first_device
+        else:
+            device_map[k] = "cuda:0"
+else:
+    first_device = device_map.get(same_device_modules[0])
+    for k in same_device_modules:
+        if k in device_map:
+            device_map[k] = first_device
+            
+model = load_checkpoint_and_dispatch(
+    model,
+    checkpoint=os.path.join(model_path, "ema.safetensors"),
+    device_map=device_map,
+    offload_buffers=True,
+    offload_folder="offload",
+    dtype=torch.bfloat16,
+    force_hooks=True,
+).eval()
+
+
+# Inferencer Preparing 
+inferencer = InterleaveInferencer(
+    model=model,
+    vae_model=vae_model,
+    tokenizer=tokenizer,
+    vae_transform=vae_transform,
+    vit_transform=vit_transform,
+    new_token_ids=new_token_ids,
+)
+
+def set_seed(seed):
+    """Set random seeds for reproducibility"""
+    if seed > 0:
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    return seed
+
 # Text to Image function with thinking option and hyperparameters
+@spaces.GPU(duration=90)
 def text_to_image(prompt, show_thinking=False, cfg_text_scale=4.0, cfg_interval=0.4, 
                  timestep_shift=3.0, num_timesteps=50, 
                  cfg_renorm_min=1.0, cfg_renorm_type="global", 
                  max_think_token_n=1024, do_sample=False, text_temperature=0.3,
                  seed=0, image_ratio="1:1"):
+    # Set seed for reproducibility
+    set_seed(seed)
+
+    if image_ratio == "1:1":
+        image_shapes = (1024, 1024)
+    elif image_ratio == "4:3":
+        image_shapes = (768, 1024)
+    elif image_ratio == "3:4":
+        image_shapes = (1024, 768) 
+    elif image_ratio == "16:9":
+        image_shapes = (576, 1024)
+    elif image_ratio == "9:16":
+        image_shapes = (1024, 576) 
+    
+    # Set hyperparameters
+    inference_hyper = dict(
+        max_think_token_n=max_think_token_n if show_thinking else 1024,
+        do_sample=do_sample if show_thinking else False,
+        text_temperature=text_temperature if show_thinking else 0.3,
+        cfg_text_scale=cfg_text_scale,
+        cfg_interval=[cfg_interval, 1.0],  # End fixed at 1.0
+        timestep_shift=timestep_shift,
+        num_timesteps=num_timesteps,
+        cfg_renorm_min=cfg_renorm_min,
+        cfg_renorm_type=cfg_renorm_type,
+        image_shapes=image_shapes,
+    )
 
-        yield None, None
+    result = {"text": "", "image": None}
+    # Call inferencer with or without think parameter based on user choice
+    for i in inferencer(text=prompt, think=show_thinking, understanding_output=False, **inference_hyper):
+        if type(i) == str:
+            result["text"] += i
+        else:
+            result["image"] = i
+
+        yield result["image"], result.get("text", None)
 
 
 # Image Understanding function with thinking option and hyperparameters
+@spaces.GPU(duration=90)
 def image_understanding(image: Image.Image, prompt: str, show_thinking=False, 
                         do_sample=False, text_temperature=0.3, max_new_tokens=512):
-    yield None
+    if image is None:
+        return "Please upload an image."
+
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    image = pil_img2rgb(image)
+    
+    # Set hyperparameters
+    inference_hyper = dict(
+        do_sample=do_sample,
+        text_temperature=text_temperature,
+        max_think_token_n=max_new_tokens, # Set max_length
+    )
+    
+    result = {"text": "", "image": None}
+    # Use show_thinking parameter to control thinking process
+    for i in inferencer(image=image, text=prompt, think=show_thinking, 
+                        understanding_output=True, **inference_hyper):
+        if type(i) == str:
+            result["text"] += i
+        else:
+            result["image"] = i
+        yield result["text"]
 
 
 # Image Editing function with thinking option and hyperparameters
+@spaces.GPU(duration=90)
 def edit_image(image: Image.Image, prompt: str, show_thinking=False, cfg_text_scale=4.0, 
               cfg_img_scale=2.0, cfg_interval=0.0, 
               timestep_shift=3.0, num_timesteps=50, cfg_renorm_min=1.0, 
               cfg_renorm_type="text_channel", max_think_token_n=1024, 
               do_sample=False, text_temperature=0.3, seed=0):
-        
-        yield (image, image), None
+    # Set seed for reproducibility
+    set_seed(seed)
+    
+    if image is None:
+        return "Please upload an image.", ""
+
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    image = pil_img2rgb(image)
+    
+    # Set hyperparameters
+    inference_hyper = dict(
+        max_think_token_n=max_think_token_n if show_thinking else 1024,
+        do_sample=do_sample if show_thinking else False,
+        text_temperature=text_temperature if show_thinking else 0.3,
+        cfg_text_scale=cfg_text_scale,
+        cfg_img_scale=cfg_img_scale,
+        cfg_interval=[cfg_interval, 1.0],  # End fixed at 1.0
+        timestep_shift=timestep_shift,
+        num_timesteps=num_timesteps,
+        cfg_renorm_min=cfg_renorm_min,
+        cfg_renorm_type=cfg_renorm_type,
+    )
+    
+    # Include thinking parameter based on user choice
+    result = {"text": "", "image": None}
+    for i in inferencer(image=image, text=prompt, think=show_thinking, understanding_output=False, **inference_hyper):
+        if type(i) == str:
+            result["text"] += i
+        else:
+            result["image"] = i
+
+        yield result["image"], result.get("text", "")
 
 # Helper function to load example images
 def load_example_image(image_path):
@@ -34,13 +267,10 @@ def load_example_image(image_path):
         print(f"Error loading example image: {e}")
         return None
 
+
 # Gradio UI 
 with gr.Blocks() as demo:
-    gr.Markdown("""
-<div>
-  <img src="https://lf3-static.bytednsdoc.com/obj/eden-cn/nuhojubrps/banner.png" alt="BAGEL" width="380"/>
-</div>
-""")
+    gr.Markdown("# 🥯 [BAGEL](https://bagel-ai.org/)")
 
     with gr.Tab("📝 Text to Image"):
         txt_input = gr.Textbox(
@@ -127,7 +357,7 @@ with gr.Blocks() as demo:
                 )
             
             with gr.Column(scale=1):
-                edit_image_output = gr.ImageSlider(label="Result")
+                edit_image_output = gr.Image(label="Result")
                 edit_thinking_output = gr.Textbox(label="Thinking Process", visible=False)
         
         with gr.Row():
@@ -233,45 +463,8 @@ with gr.Blocks() as demo:
             outputs=txt_output
         )
 
-    gr.Markdown("""
-<div style="display: flex; justify-content: flex-start; flex-wrap: wrap; gap: 10px;">
-  <a href="https://bagel-ai.org/">
-    <img
-      src="https://img.shields.io/badge/BAGEL-Website-0A66C2?logo=safari&logoColor=white"
-      alt="BAGEL Website"
-    />
-  </a>
-  <a href="https://arxiv.org/abs/2505.14683">
-    <img
-      src="https://img.shields.io/badge/BAGEL-Paper-red?logo=arxiv&logoColor=red"
-      alt="BAGEL Paper on arXiv"
-    />
-  </a>
-  <a href="https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT">
-    <img 
-        src="https://img.shields.io/badge/BAGEL-Hugging%20Face-orange?logo=huggingface&logoColor=yellow" 
-        alt="BAGEL on Hugging Face"
-    />
-  </a>
-  <a href="https://demo.bagel-ai.org/">
-    <img
-      src="https://img.shields.io/badge/BAGEL-Demo-blue?logo=googleplay&logoColor=blue"
-      alt="BAGEL Demo"
-    />
-  </a>
-  <a href="https://discord.gg/Z836xxzy">
-    <img
-      src="https://img.shields.io/badge/BAGEL-Discord-5865F2?logo=discord&logoColor=purple"
-      alt="BAGEL Discord"
-    />
-  </a>
-  <a href="mailto:bagel@bytedance.com">
-    <img
-      src="https://img.shields.io/badge/BAGEL-Email-D14836?logo=gmail&logoColor=red"
-      alt="BAGEL Email"
-    />
-  </a>
-</div>
-""")
-
-demo.launch()
+    gr.Markdown(
+        "🌐[Website](https://bagel-ai.org/)&nbsp;&nbsp;📄[Report](https://arxiv.org/abs/2505.14683)&nbsp;&nbsp;🤗[Model](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT)&nbsp;&nbsp;🚀[Demo](https://demo.bagel-ai.org/)&nbsp;&nbsp;💬[Discord](https://discord.gg/Z836xxzy)&nbsp;&nbsp;📧[Contact](mailto:bagel@bytedance.com)"
+    )
+
+demo.launch(share=True)