init

app.py

@@ -1,264 +1,30 @@
-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,
-    )
 
-    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):
-        print(type(i))
-        if type(i) == str:
-            result["text"] += i
-        else:
-            result["image"] = i
-
-        yield result["image"], result.get("text", None)
+        yield None, 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):
-    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"]
+    yield None
 
 
 # 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):
-    # 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", "")
+        yield None, None
 
 # Helper function to load example images
 def load_example_image(image_path):
@@ -268,7 +34,6 @@ def load_example_image(image_path):
         print(f"Error loading example image: {e}")
         return None
 
-
 # Gradio UI 
 with gr.Blocks() as demo:
     gr.Markdown("""