Update app.py

app.py

@@ -1,326 +1,123 @@
-import gradio as gr
-import os
-import torch
-import spaces
-
-from llava import conversation as conversation_lib
-from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
-from llava.conversation import conv_templates, SeparatorStyle
-from llava.model.builder import load_pretrained_model
-from llava.utils import disable_torch_init
-from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, process_images
+# gradio_blip3o_next_min.py
+import time
+from dataclasses import dataclass
 
+import torch
 from PIL import Image
-import argparse
-
-from transformers import TextIteratorStreamer
-from threading import Thread
-
-import subprocess
-# Install flash attention, skipping CUDA build if necessary
-subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
-)
-
-
-# os.environ['GRADIO_TEMP_DIR'] = './gradio_tmp'
-no_change_btn = gr.Button()
-enable_btn = gr.Button(interactive=True)
-disable_btn = gr.Button(interactive=False)
-
-argparser = argparse.ArgumentParser()
-argparser.add_argument("--server_name", default="0.0.0.0", type=str)
-argparser.add_argument("--port", default="6324", type=str)
-argparser.add_argument("--model-path", default="jiuhai/florence-phi-ms", type=str)
-argparser.add_argument("--model-base", type=str, default=None)
-argparser.add_argument("--num-gpus", type=int, default=1)
-argparser.add_argument("--conv-mode", type=str, default="llama3")
-argparser.add_argument("--temperature", type=float, default=0.2)
-argparser.add_argument("--max-new-tokens", type=int, default=512)
-argparser.add_argument("--num_frames", type=int, default=16)
-argparser.add_argument("--load-8bit", action="store_true")
-argparser.add_argument("--load-4bit", action="store_true")
-argparser.add_argument("--debug", action="store_true")
-
-args = argparser.parse_args()
-model_path = args.model_path
-conv_mode = args.conv_mode
-filt_invalid="cut"
-model_name = get_model_name_from_path(args.model_path)
-model_kwargs = {
-    "trust_remote_code": True,
-    "torch_dtype": torch.bfloat16,
-    "attn_implementation": "eager"
-}
-tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, device_map="cuda:0", **model_kwargs)
-our_chatbot = None
-
-def upvote_last_response(state):
-    return ("",) + (disable_btn,) * 3
-
-
-def downvote_last_response(state):
-    return ("",) + (disable_btn,) * 3
-
-
-def flag_last_response(state):
-    return ("",) + (disable_btn,) * 3
-
-def clear_history():
-    state =conv_templates[conv_mode].copy()
-    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
-
-def add_text(state, imagebox, textbox, image_process_mode):
-    if state is None:
-        state = conv_templates[conv_mode].copy()
-
-    if imagebox is not None:
-        textbox = DEFAULT_IMAGE_TOKEN + '\n' + textbox
-        image = Image.open(imagebox).convert('RGB')
-
-    if imagebox is not None:
-        textbox = (textbox, image, image_process_mode)
-
-    state.append_message(state.roles[0], textbox)
-    state.append_message(state.roles[1], None)
-
-    yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
-
-def delete_text(state, image_process_mode):
-    state.messages[-1][-1] = None
-    prev_human_msg = state.messages[-2]
-    if type(prev_human_msg[1]) in (tuple, list):
-        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
-    yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
-
-def regenerate(state, image_process_mode):
-    state.messages[-1][-1] = None
-    prev_human_msg = state.messages[-2]
-    if type(prev_human_msg[1]) in (tuple, list):
-        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
-    state.skip_next = False
-    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
-
-@spaces.GPU
-def generate(state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens):
-    prompt = state.get_prompt()
-    images = state.get_images(return_pil=True)
-    #prompt, image_args = process_image(prompt, images)
-
-    ori_prompt = prompt
-    num_image_tokens = 0
-
-    if images is not None and len(images) > 0:
-        if len(images) > 0:
-            if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
-                raise ValueError("Number of images does not match number of <image> tokens in prompt")
-            
-            #images = [load_image_from_base64(image) for image in images]
-            image_sizes = [image.size for image in images]
-            images = process_images(images, image_processor, model.config)
-
-            if type(images) is list:
-                images = [image.to(model.device, dtype=torch.float16) for image in images]
-            else:
-                images = images.to(model.device, dtype=torch.float16)
-        else:
-            images = None
-            image_sizes = None
-        image_args = {"images": images, "image_sizes": image_sizes}
-    else:
-        images = None
-        image_args = {}
-
-    max_context_length = getattr(model.config, 'max_position_embeddings', 2048)
-    max_new_tokens = 512
-    do_sample = True if temperature > 0.001 else False
-    stop_str = state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2
-
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
-
-    max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens)
-
-    if max_new_tokens < 1:
-        # yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"
-        return
-
-    thread = Thread(target=model.generate, kwargs=dict(
-        inputs=input_ids,
-        do_sample=do_sample,
-        temperature=temperature,
-        top_p=top_p,
-        max_new_tokens=max_new_tokens,
-        streamer=streamer,
-        use_cache=True,
-        pad_token_id=tokenizer.eos_token_id,
-        eos_token_id=[32007],
-        **image_args
-    ))
-    thread.start()
-    generated_text = ''
-    for new_text in streamer:
-        new_text = new_text.replace('<|end|>', "")
-        generated_text += new_text
-        if generated_text.endswith(stop_str):
-            generated_text = generated_text[:-len(stop_str)]
-        state.messages[-1][-1] = generated_text
-        yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
-    
-    yield (state, state.to_gradio_chatbot(), "", None) + (enable_btn,) * 5
-    
-    torch.cuda.empty_cache()
-
-txt = gr.Textbox(
-    scale=4,
-    show_label=False,
-    placeholder="Enter text and press enter.",
-    container=False,
-)
-
-
-# title_markdown = ("""
-# # llava: Exploring The Design Space for Multimodal LLMs with Mixture of Encoders
-# [[Code](https://github.com/NVlabs/llava)] [[Model](https://huggingface.co/NVllava)] | 📚 [[Arxiv](https://arxiv.org/pdf/2408.15998)]]
-# """)
-
-title_markdown = ("""
-# Florence-phi
-""")
-
-tos_markdown = ("""
-### Terms of use
-By using this service, users are required to agree to the following terms:
-The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
-Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
-For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
-""")
-
-
-learn_more_markdown = ("""
-### License
-The service is a research preview intended for non-commercial use only, subject to the. Please contact us if you find any potential violation.
-""")
-
-block_css = """
-#buttons button {
-    min-width: min(120px,100%);
-}
-"""
-
-textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
-with gr.Blocks(title="llava", theme=gr.themes.Default(), css=block_css) as demo:
-    state = gr.State()
-
-    gr.Markdown(title_markdown)
+from transformers import AutoTokenizer
+from blip3o.model import *
+import gradio as gr
+from huggingface_hub import snapshot_download
+
+
+# -----------------------------
+# Minimal config and runner
+# -----------------------------
+@dataclass
+class T2IConfig:
+    device: str = "cuda:0"
+    dtype: torch.dtype = torch.bfloat16
+    # fixed generation config (no UI controls)
+    scale: int = 0
+    seq_len: int = 729
+    top_p: float = 0.95
+    top_k: int = 1200
+
+
+class TextToImageInference:
+    def __init__(self, config: T2IConfig):
+        self.config = config
+        self.device = torch.device(config.device)
+        self._load_models()
+
+    def _load_models(self):
+        model_path = snapshot_download(repo_id='BLIP3o/BLIP3o-NEXT-GRPO-Geneval-3B')
+        self.model = blip3oQwenForInferenceLM.from_pretrained(
+            model_path, torch_dtype=self.config.dtype
+        ).to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        if hasattr(self.tokenizer, "padding_side"):
+            self.tokenizer.padding_side = "left"
+
+    @torch.inference_mode()
+    def generate_image(self, prompt: str) -> Image.Image:
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {
+                "role": "user",
+                "content": f"Please generate image based on the following caption: {prompt}",
+            },
+        ]
+        input_text = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        input_text += f"<im_start><S{self.config.scale}>"
+
+        inputs = self.tokenizer(
+            [input_text], return_tensors="pt", padding=True, truncation=True
+        )
+
+        _, images = self.model.generate_images(
+            inputs.input_ids.to(self.device),
+            inputs.attention_mask.to(self.device),
+            max_new_tokens=self.config.seq_len,
+            do_sample=True,
+            top_p=self.config.top_p,
+            top_k=self.config.top_k,
+        )
+        return images[0]
+
+
+# Try loading once at startup for simplicity
+LOAD_ERROR = None
+inference = None
+try:
+    inference = TextToImageInference(T2IConfig())
+except Exception as e:
+    LOAD_ERROR = f"❌ Failed to load model: {e}"
+
+
+def run_generate(prompt, progress=gr.Progress(track_tqdm=True)):
+    t0 = time.time()
+    if LOAD_ERROR:
+        return None, LOAD_ERROR
+    if not prompt or not prompt.strip():
+        return None, "⚠️ Please enter a prompt."
+
+    try:
+        img = inference.generate_image(prompt.strip())
+        return img, f"✅ Done in {time.time() - t0:.2f}s."
+    except torch.cuda.OutOfMemoryError:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return None, "❌ CUDA OOM. Try reducing other GPU workloads."
+    except Exception as e:
+        return None, f"❌ Error: {e}"
+
+
+with gr.Blocks(title="BLIP3o-NEXT-GRPO-Geneval — Text ➜ Image") as demo:
+    gr.Markdown("# BLIP3o-NEXT-GRPO-Geneval — Text ➜ Image")
 
     with gr.Row():
         with gr.Column(scale=3):
-            imagebox = gr.Image(label="Input Image", type="filepath")
-            image_process_mode = gr.Radio(
-                ["Crop", "Resize", "Pad", "Default"],
-                value="Default",
-                label="Preprocess for non-square image", visible=False)
-    
-            cur_dir = os.path.dirname(os.path.abspath(__file__))
-            # gr.Examples(examples=[
-            #     [f"{cur_dir}/assets/health-insurance.png", "Under which circumstances do I need to be enrolled in mandatory health insurance if I am an international student?"],
-            #     [f"{cur_dir}/assets/leasing-apartment.png", "I don't have any 3rd party renter's insurance now. Do I need to get one for myself?"],
-            #     [f"{cur_dir}/assets/nvidia.jpeg", "Who is the person in the middle?"],
-            #     [f"{cur_dir}/assets/animal-compare.png", "Are these two pictures showing the same kind of animal?"],
-            #     [f"{cur_dir}/assets/georgia-tech.jpeg", "Where is this photo taken?"]
-            # ], inputs=[imagebox, textbox], cache_examples=False)
-
-            gr.Examples(examples=[
-                [f"{cur_dir}/assets/animal-compare.png", "Are these two pictures showing the same kind of animal?"]
-            ], inputs=[imagebox, textbox], cache_examples=False)
-
-            with gr.Accordion("Parameters", open=False) as parameter_row:
-                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, interactive=True, label="Temperature",)
-                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P",)
-                max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)
-
-        with gr.Column(scale=8):
-            chatbot = gr.Chatbot(
-                elem_id="chatbot",
-                label="llava Chatbot",
-                height=650,
-                layout="panel",
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Describe the image you want to generate...",
+                lines=4,
             )
-            with gr.Row():
-                with gr.Column(scale=8):
-                    textbox.render()
-                with gr.Column(scale=1, min_width=50):
-                    submit_btn = gr.Button(value="Send", variant="primary")
-            with gr.Row(elem_id="buttons") as button_row:
-                upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
-                downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
-                flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
-                #stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
-                regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
-                clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
-
-    gr.Markdown(tos_markdown)
-    gr.Markdown(learn_more_markdown)
-    url_params = gr.JSON(visible=False)
-
-    # Register listeners
-    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
-    upvote_btn.click(
-        upvote_last_response,
-        [state],
-        [textbox, upvote_btn, downvote_btn, flag_btn]
-    )
-    downvote_btn.click(
-        downvote_last_response,
-        [state],
-        [textbox, upvote_btn, downvote_btn, flag_btn]
-    )
-    flag_btn.click(
-        flag_last_response,
-        [state],
-        [textbox, upvote_btn, downvote_btn, flag_btn]
-    )
-
-    clear_btn.click(
-        clear_history,
-        None,
-        [state, chatbot, textbox, imagebox] + btn_list,
-        queue=False
-    )
-
-    regenerate_btn.click(
-        delete_text,
-        [state, image_process_mode],
-        [state, chatbot, textbox, imagebox] + btn_list,
-    ).then(
-        generate,
-        [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens],
-        [state, chatbot, textbox, imagebox] + btn_list,
-    )
-    textbox.submit(
-        add_text,
-        [state, imagebox, textbox, image_process_mode],
-        [state, chatbot, textbox, imagebox] + btn_list,
-    ).then(
-        generate,
-        [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens],
-        [state, chatbot, textbox, imagebox] + btn_list,
-    )
-
-    submit_btn.click(
-        add_text,
-        [state, imagebox, textbox, image_process_mode],
-        [state, chatbot, textbox, imagebox] + btn_list,
-    ).then(
-        generate,
-        [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens],
-        [state, chatbot, textbox, imagebox] + btn_list,
+            run_btn = gr.Button("Generate", variant="primary")
+
+        with gr.Column(scale=4):
+            out_img = gr.Image(label="Generated Image", format="png")
+            status = gr.Markdown("")
+
+    run_btn.click(
+        fn=run_generate,
+        inputs=[prompt],
+        outputs=[out_img, status],
+        queue=True,
+        api_name="generate",
     )
 
-demo.queue(
-    status_update_rate=10,
-    api_open=False
-).launch()
+if __name__ == "__main__":
+    demo.queue().launch(share=True)