Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| import numpy as np | |
| import random | |
| import spaces #[uncomment to use ZeroGPU] | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from transformers.generation import GenerationConfig | |
| device = "cuda" | |
| torch_dtype = torch.bfloat16 | |
| model_name_or_path = "X-Omni/X-Omni-Zh" | |
| flux_model_name_or_path = "zhangxiaosong18/FLUX.1-dev-VAE" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name_or_path, | |
| torch_dtype=torch_dtype, | |
| trust_remote_code=True, | |
| ).cuda() | |
| model.init_vision(flux_model_name_or_path) | |
| model.set_generation_mode('image') | |
| model.eval() | |
| #[uncomment to use ZeroGPU] | |
| def generate_image( | |
| image_prompt, | |
| image_size, | |
| top_p, | |
| min_p, | |
| seed, | |
| ): | |
| image_prompt = image_prompt.strip() | |
| image_size = tuple(map(int, image_size.split('x'))) | |
| token_h, token_w = image_size[0] // 16, image_size[1] // 16 | |
| image_prefix = f'<SOM>{token_h} {token_w}<IMAGE>' | |
| generation_config = GenerationConfig( | |
| max_new_tokens=token_h * token_w, | |
| do_sample=True, | |
| temperature=1.0, | |
| min_p=min_p, | |
| top_p=top_p, | |
| guidance_scale=1.0, | |
| suppress_tokens=tokenizer.convert_tokens_to_ids(model.config.mm_special_tokens), | |
| ) | |
| tokens = tokenizer( | |
| [image_prompt + image_prefix], | |
| return_tensors='pt', | |
| padding='longest', | |
| padding_side='left', | |
| ) | |
| input_ids = tokens.input_ids.cuda() | |
| attention_mask = tokens.attention_mask.cuda() | |
| torch.manual_seed(seed) | |
| tokens = model.generate( | |
| inputs=input_ids, | |
| attention_mask=attention_mask, | |
| generation_config=generation_config, | |
| ) | |
| torch.manual_seed(seed) | |
| _, images = model.mmdecode(tokenizer, tokens[0], skip_special_tokens=False) | |
| return images[0] | |
| examples = [ | |
| ''' | |
| 生成一张雪中的紫禁城全景封面图，作为北京冬季旅游指南的主题。画面以近景构图展现建筑，红墙金瓦被皑皑白雪覆盖，朱红色宫墙，金黄色瓦片与洁白雪色形成强烈对比，琉璃瓦顶的积雪在阳光下折射出晶莹光泽。前景一枝腊梅花正在盛开，背景为灰蓝色冬日天空，飘落细雪，远处角楼轮廓若隐若现，增添朦胧诗意感。图片上有标题“雪落北平·穿越600年”，另有副标题“北京古建筑雪景深度游”。文字艺术感极强，与图片良好融合起来 | |
| '''.strip(), | |
| ''' | |
| 画面的中心摆放着一个复古花瓶，瓶身主体为浓郁的蓝色，这种蓝色深邃而典雅，仿佛带着岁月的沉淀。花瓶设计极具复古风格，瓶颈处环绕着细致的金色雕花，宛如华丽的项链点缀其上；瓶身绘制着精美的花卉图案，笔触细腻，色彩过渡自然，展现出极高的工艺水准，整体彰显出优雅的古典韵味。花瓶放置在深色木质的圆桌上，旁边搭配了一束新鲜绽放的百合花，为画面增添了几分生机与活力。背景是一幅淡蓝色的壁纸，上面有着若隐若现的花纹，营造出一种静谧而温馨的氛围。图片中的文字信息十分醒目。“家居美学盛典”位于顶部中央，字体较大，在视觉上十分突出，吸引观众的目光；左下角写着“下单直降 100”，下方紧跟数字“399”，强调了价格优惠；右下角有“限量抢购 速来咨询”的提示，引导观众进一步咨询；最底部中央，“前 50 名买一送一”的字样突出促销活动的紧迫性和吸引力。这些文字信息通过巧妙的颜色、大小和背景设计，在空间布局上层次分明，重点突出，有效地引导观众关注促销信息和价格优势。 | |
| '''.strip(), | |
| ] | |
| examples = [[prompt, '1152x1152', 1.0, 0.03, 0] for prompt in examples] | |
| css = """ | |
| .app { | |
| max-width: 800px !important; | |
| margin: 0 auto !important; | |
| } | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| gr.HTML(''' | |
| <h1 style="text-align:center">🎨X-Omni: Reinforcement Learning Makes Discrete Autoregressive Image Generative Models Great Again</h1> | |
| <h3 style="text-align:center">Model: <a href="https://huggingface.co/X-Omni/X-Omni-Zh">X-Omni-Zh</a> (support Chinese text rendering)</h3> | |
| <p align="center"> | |
| <a href="https://x-omni-team.github.io">🏠 Project Page</a> | | |
| <a href="https://arxiv.org/pdf/2507.22058">📄 Paper</a> | | |
| <a href="https://github.com/X-Omni-Team/X-Omni">💻 Code</a> | | |
| <a href="https://huggingface.co/collections/X-Omni/x-omni-models-6888aadcc54baad7997d7982">🤗 HuggingFace Model</a> | |
| </p> | |
| '''.strip()) | |
| with gr.Row(): | |
| textbox = gr.Textbox(lines=2, placeholder='text prompt for image generation', show_label=False) | |
| image = gr.Image(show_label=False, type='pil') | |
| with gr.Row(): | |
| button = gr.Button("Generate", variant="primary") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| image_size = gr.Dropdown(label="Image Size", choices=["1152x1152", "1152x768", "768x1152"], value="1152x1152") | |
| top_p = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=1.0, step=0.01) | |
| min_p = gr.Slider(label="Min P", minimum=0.0, maximum=1.0, value=0.03, step=0.01) | |
| seed_input = gr.Number(label="Seed", value=0, precision=0) | |
| with gr.Row(): | |
| gr.Examples(examples=examples, inputs=(textbox, image_size, top_p, min_p, seed_input), outputs=image, fn=generate_image, cache_examples=False, run_on_click=True) | |
| button.click( | |
| generate_image, | |
| inputs=(textbox, image_size, top_p, min_p, seed_input), | |
| outputs=image, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(ssr_mode=False) | |