Create app.py

app.py

@@ -0,0 +1,117 @@
+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()
+
+@spaces.GPU(duration=220) #[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://x-omni-team.github.io">📄 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)