import os import re from PIL import Image import spaces import gradio as gr import uuid import argparse from huggingface_hub import login, snapshot_download import torch from dreamomni2.pipeline_dreamomni2 import DreamOmni2Pipeline from diffusers.utils import load_image from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor from utils.vprocess import process_vision_info, resizeinput def extract_gen_content(text): text = text[6:-7] return text def _load_model_processor(): device = "cuda" if torch.cuda.is_available() else "cpu" local_dir = snapshot_download( repo_id="xiabs/DreamOmni2", revision="main", allow_patterns=["vlm-model/**", "gen_lora/**"], ) vlm_dir = os.path.join(local_dir, 'vlm-model') lora_dir = os.path.join(local_dir, 'gen_lora') print(f"Loading models from vlm_path: {vlm_dir}, gen_lora_path: {lora_dir}") pipe = DreamOmni2Pipeline.from_pretrained( "black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16 ).to(device) pipe.load_lora_weights(lora_dir, adapter_name="generation") pipe.set_adapters(["generation"], adapter_weights=[1]) vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained( vlm_dir, torch_dtype="bfloat16" ).to(device) processor = AutoProcessor.from_pretrained(vlm_dir) return vlm_model, processor, pipe def _launch_demo(vlm_model, processor, pipe): @spaces.GPU(duration=90) def infer_vlm(input_img_path, input_instruction, prefix): if not vlm_model or not processor: raise gr.Error("VLM Model not loaded. Cannot process prompt.") tp = [] for path in input_img_path: tp.append({"type": "image", "image": path}) tp.append({"type": "text", "text": input_instruction + prefix}) messages = [{"role": "user", "content": tp}] text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) image_inputs, video_inputs = process_vision_info(messages) inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt") inputs = inputs.to(device=vlm_model.device) generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=4096) generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)] output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False) return output_text[0] PREFERRED_KONTEXT_RESOLUTIONS = [ (672, 1568), (688, 1504), (720, 1456), (752, 1392), (800, 1328), (832, 1248), (880, 1184), (944, 1104), (1024, 1024), (1104, 944), (1184, 880), (1248, 832), (1328, 800), (1392, 752), (1456, 720), (1504, 688), (1568, 672), ] def find_closest_resolution(width, height, preferred_resolutions): input_ratio = width / height closest_resolution = min( preferred_resolutions, key=lambda res: abs((res[0] / res[1]) - input_ratio) ) return closest_resolution @spaces.GPU(duration=90) def perform_generation(input_img_paths, input_instruction, output_path, height=1024, width=1024): prefix = " It is generation task." source_imgs = [] for path in input_img_paths: img = load_image(path) # source_imgs.append(img) source_imgs.append(resizeinput(img)) prompt = infer_vlm(input_img_paths, input_instruction, prefix) prompt = extract_gen_content(prompt) print(f"Generated Prompt for VLM: {prompt}") image = pipe( images=source_imgs, height=height, width=width, prompt=prompt, num_inference_steps=30, guidance_scale=3.5, ).images[0] image.save(output_path) print(f"Generation result saved to {output_path}") @spaces.GPU(duration=90) def process_request(image_file_1, image_file_2, instruction): # debugpy.listen(5678) # print("Waiting for debugger attach...") # debugpy.wait_for_client() if not image_file_1 or not image_file_2: raise gr.Error("Please upload both images.") if not instruction: raise gr.Error("Please provide an instruction.") if not pipe or not vlm_model: raise gr.Error("Models not loaded. Check the console for errors.") output_path = f"/tmp/{uuid.uuid4()}.png" input_img_paths = [image_file_1, image_file_2] # List of file paths from the two gr.File inputs perform_generation(input_img_paths, instruction, output_path) return output_path css = """ .text-center { text-align: center; } .result-img img { max-height: 60vh !important; min-height: 30vh !important; width: auto !important; object-fit: contain; } .input-img img { max-height: 30vh !important; width: auto !important; object-fit: contain; } """ with gr.Blocks(theme=gr.themes.Soft(), title="DreamOmni2", css=css) as demo: gr.HTML( """

DreamOmni2: Multimodal Image Generation and Editing

""" ) gr.Markdown( "Upload two images, provide an instruction, and click 'Run'.", elem_classes="text-center" ) with gr.Row(): with gr.Column(scale=2): gr.Markdown("⬆️ Upload images. Click or drag to upload.") with gr.Row(): image_uploader_1 = gr.Image( label="Img 1", type="filepath", interactive=True, elem_classes="input-img", ) image_uploader_2 = gr.Image( label="Img 2", type="filepath", interactive=True, elem_classes="input-img", ) instruction_text = gr.Textbox( label="Instruction", lines=2, placeholder="Input your instruction for generation or editing here...", ) run_button = gr.Button("Run", variant="primary") with gr.Column(scale=2): gr.Markdown("🖼️ **Generation Mode**: Create new scenes from reference images.\n\n" "Tip: If the result is not what you expect, try clicking **Run** again. ") output_image = gr.Image( label="Result", type="filepath", elem_classes="result-img", ) # --- Examples --- gr.Markdown("## Examples") gr.Examples( label="Generation Examples", examples=[ [ "example_input/gen_tests/img1.jpg", "example_input/gen_tests/img2.jpg", "In the scene, the character from the first image stands on the left, and the character from the second image stands on the right. They are shaking hands against the backdrop of a spaceship interior.", "example_input/gen_tests/gen_res.png" ] ], inputs=[image_uploader_1, image_uploader_2, instruction_text, output_image], cache_examples=False, ) run_button.click( fn=process_request, inputs=[image_uploader_1, image_uploader_2, instruction_text], outputs=output_image ) demo.launch() if __name__ == "__main__": vlm_model, processor, pipe = _load_model_processor() _launch_demo(vlm_model, processor, pipe)