app.py

import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import os.path as osp
import time
import hashlib
import argparse
import random
from pathlib import Path
from typing import List, Dict, Optional
from dataclasses import dataclass

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageEnhance
from torchvision.transforms.functional import to_tensor
from transformers import AutoTokenizer, T5EncoderModel, T5TokenizerFast
from huggingface_hub import hf_hub_download
import gradio as gr
import spaces

from models.infinity import Infinity
from models.basic import *
from utils.dynamic_resolution import dynamic_resolution_h_w
from gradio_client import Client

# Performance optimizations
torch._dynamo.config.cache_size_limit = 64
torch.backends.cudnn.benchmark = True  # Enable cudnn auto-tuner
client = Client("Qwen/Qwen2.5-72B-Instruct")

@dataclass
class ModelConfig:
    """Configuration for Infinity model."""
    depth: int
    embed_dim: int
    num_heads: int
    drop_path_rate: float = 0.1
    mlp_ratio: float = 4.0
    block_chunks: int = 8

    @classmethod
    def from_type(cls, model_type: str) -> 'ModelConfig':
        """Create model config from predefined types."""
        configs = {
            'infinity_2b': dict(depth=32, embed_dim=2048, num_heads=2048//128),
            'infinity_layer12': dict(depth=12, embed_dim=768, num_heads=8),
            'infinity_layer16': dict(depth=16, embed_dim=1152, num_heads=12),
            'infinity_layer24': dict(depth=24, embed_dim=1536, num_heads=16),
            'infinity_layer32': dict(depth=32, embed_dim=2080, num_heads=20),
            'infinity_layer40': dict(depth=40, embed_dim=2688, num_heads=24),
            'infinity_layer48': dict(depth=48, embed_dim=3360, num_heads=28),
        }
        if model_type not in configs:
            raise ValueError(f"Unknown model type: {model_type}")
        return cls(**configs[model_type])

    def to_dict(self) -> Dict:
        """Convert config to dictionary."""
        return {
            'depth': self.depth,
            'embed_dim': self.embed_dim,
            'num_heads': self.num_heads,
            'drop_path_rate': self.drop_path_rate,
            'mlp_ratio': self.mlp_ratio,
            'block_chunks': self.block_chunks
        }

# Global device configuration
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Define a function to download weights if not present
def download_infinity_weights(weights_path):
    try:
        model_file = weights_path / 'infinity_2b_reg.pth'
        if not model_file.exists():
            hf_hub_download(repo_id="FoundationVision/Infinity", filename="infinity_2b_reg.pth", local_dir=str(weights_path))
        
        vae_file = weights_path / 'infinity_vae_d32reg.pth'
        if not vae_file.exists():
            hf_hub_download(repo_id="FoundationVision/Infinity", filename="infinity_vae_d32reg.pth", local_dir=str(weights_path))
        
    except Exception as e:
        print(f"Error downloading weights: {e}")

def extract_key_val(text):
    pattern = r'<(.+?):(.+?)>'
    matches = re.findall(pattern, text)
    key_val = {}
    for match in matches:
        key_val[match[0]] = match[1].lstrip()
    return key_val

def encode_prompt(text_tokenizer, text_encoder, prompt, enable_positive_prompt=False):
    if enable_positive_prompt:
        print(f'before positive_prompt aug: {prompt}')
        prompt = aug_with_positive_prompt(prompt)
        print(f'after positive_prompt aug: {prompt}')
    print(f'prompt={prompt}')
    captions = [prompt]
    tokens = text_tokenizer(text=captions, max_length=512, padding='max_length', truncation=True, return_tensors='pt')  # todo: put this into dataset
    input_ids = tokens.input_ids.cuda(non_blocking=True)
    mask = tokens.attention_mask.cuda(non_blocking=True)
    text_features = text_encoder(input_ids=input_ids, attention_mask=mask)['last_hidden_state'].float()
    lens: List[int] = mask.sum(dim=-1).tolist()
    cu_seqlens_k = F.pad(mask.sum(dim=-1).to(dtype=torch.int32).cumsum_(0), (1, 0))
    Ltext = max(lens)    
    kv_compact = []
    for len_i, feat_i in zip(lens, text_features.unbind(0)):
        kv_compact.append(feat_i[:len_i])
    kv_compact = torch.cat(kv_compact, dim=0)
    text_cond_tuple = (kv_compact, lens, cu_seqlens_k, Ltext)
    return text_cond_tuple

def aug_with_positive_prompt(prompt):
    for key in ['man', 'woman', 'men', 'women', 'boy', 'girl', 'child', 'person', 'human', 'adult', 'teenager', 'employee', 
                'employer', 'worker', 'mother', 'father', 'sister', 'brother', 'grandmother', 'grandfather', 'son', 'daughter']:
        if key in prompt:
            prompt = prompt + '. very smooth faces, good looking faces, face to the camera, perfect facial features'
            break
    return prompt

def enhance_image(image):
    for t in range(1):
        contrast_image = image.copy()
        contrast_enhancer = ImageEnhance.Contrast(contrast_image)
        contrast_image = contrast_enhancer.enhance(1.05)  # 增强对比度
        color_image = contrast_image.copy()
        color_enhancer = ImageEnhance.Color(color_image)
        color_image = color_enhancer.enhance(1.05)  # 增强饱和度
    return color_image

def get_prompt_id(prompt):
    md5 = hashlib.md5()
    md5.update(prompt.encode('utf-8'))
    prompt_id = md5.hexdigest()
    return prompt_id

def save_slim_model(infinity_model_path, save_file=None, device='cpu', key='gpt_fsdp'):
    print('[Save slim model]')
    full_ckpt = torch.load(infinity_model_path, map_location=device)
    infinity_slim = full_ckpt['trainer'][key]
    # ema_state_dict = cpu_d['trainer'].get('gpt_ema_fsdp', state_dict)
    if not save_file:
        save_file = osp.splitext(infinity_model_path)[0] + '-slim.pth'
    print(f'Save to {save_file}')
    torch.save(infinity_slim, save_file)
    print('[Save slim model] done')
    return save_file

def load_tokenizer(t5_path =''):
    print(f'[Loading tokenizer and text encoder]')
    text_tokenizer: T5TokenizerFast = AutoTokenizer.from_pretrained(t5_path, revision=None, legacy=True)
    text_tokenizer.model_max_length = 512
    text_encoder: T5EncoderModel = T5EncoderModel.from_pretrained(t5_path, torch_dtype=torch.float16)
    text_encoder.to(DEVICE)
    text_encoder.eval()
    text_encoder.requires_grad_(False)
    return text_tokenizer, text_encoder

def load_infinity(
    rope2d_each_sa_layer, 
    rope2d_normalized_by_hw, 
    use_scale_schedule_embedding, 
    pn, 
    use_bit_label, 
    add_lvl_embeding_only_first_block, 
    model_path='', 
    scale_schedule=None, 
    vae=None, 
    model_kwargs=None,
    text_channels=2048,
    apply_spatial_patchify=0,
    use_flex_attn=False,
    bf16=False,
):
    print(f'[Loading Infinity]')
    
    # Set autocast dtype based on bf16 and device support
    if bf16 and DEVICE.type == 'cuda' and torch.cuda.is_bf16_supported():
        autocast_dtype = torch.bfloat16
    else:
        autocast_dtype = torch.float32
        bf16 = False  # Disable bf16 if not supported

    text_maxlen = 512
    torch.cuda.empty_cache()

    with torch.amp.autocast(device_type=DEVICE.type, dtype=autocast_dtype), torch.no_grad():
        infinity_test: Infinity = Infinity(
            vae_local=vae, text_channels=text_channels, text_maxlen=text_maxlen,
            shared_aln=True, raw_scale_schedule=scale_schedule,
            checkpointing='full-block',
            customized_flash_attn=False,
            fused_norm=True,
            pad_to_multiplier=128,
            use_flex_attn=use_flex_attn,
            add_lvl_embeding_only_first_block=add_lvl_embeding_only_first_block,
            use_bit_label=use_bit_label,
            rope2d_each_sa_layer=rope2d_each_sa_layer,
            rope2d_normalized_by_hw=rope2d_normalized_by_hw,
            pn=pn,
            apply_spatial_patchify=apply_spatial_patchify,
            inference_mode=True,
            train_h_div_w_list=[1.0],
            **model_kwargs,
        ).to(DEVICE)

        print(f'[you selected Infinity with {model_kwargs=}] model size: {sum(p.numel() for p in infinity_test.parameters())/1e9:.2f}B, bf16={bf16}')

        if bf16:
            for block in infinity_test.unregistered_blocks:
                block.bfloat16()

        infinity_test.eval()
        infinity_test.requires_grad_(False)

        print(f'[Load Infinity weights]')
        state_dict = torch.load(model_path, map_location=DEVICE)
        print(infinity_test.load_state_dict(state_dict))
        
        # Initialize random number generator
        infinity_test.rng = torch.Generator(device=DEVICE)
        
        return infinity_test

def transform(pil_img, tgt_h, tgt_w):
    width, height = pil_img.size
    if width / height <= tgt_w / tgt_h:
        resized_width = tgt_w
        resized_height = int(tgt_w / (width / height))
    else:
        resized_height = tgt_h
        resized_width = int((width / height) * tgt_h)
    pil_img = pil_img.resize((resized_width, resized_height), resample=PImage.LANCZOS)
    # crop the center out
    arr = np.array(pil_img)
    crop_y = (arr.shape[0] - tgt_h) // 2
    crop_x = (arr.shape[1] - tgt_w) // 2
    im = to_tensor(arr[crop_y: crop_y + tgt_h, crop_x: crop_x + tgt_w])
    return im.add(im).add_(-1)

def joint_vi_vae_encode_decode(vae, image_path, scale_schedule, device, tgt_h, tgt_w):
    pil_image = Image.open(image_path).convert('RGB')
    inp = transform(pil_image, tgt_h, tgt_w)
    inp = inp.unsqueeze(0).to(device)
    scale_schedule = [(item[0], item[1], item[2]) for item in scale_schedule]
    t1 = time.time()
    h, z, _, all_bit_indices, _, infinity_input = vae.encode(inp, scale_schedule=scale_schedule)
    t2 = time.time()
    recons_img = vae.decode(z)[0]
    if len(recons_img.shape) == 4:
        recons_img = recons_img.squeeze(1)
    print(f'recons: z.shape: {z.shape}, recons_img shape: {recons_img.shape}')
    t3 = time.time()
    print(f'vae encode takes {t2-t1:.2f}s, decode takes {t3-t2:.2f}s')
    recons_img = (recons_img + 1) / 2
    recons_img = recons_img.permute(1, 2, 0).mul_(255).cpu().numpy().astype(np.uint8)
    gt_img = (inp[0] + 1) / 2
    gt_img = gt_img.permute(1, 2, 0).mul_(255).cpu().numpy().astype(np.uint8)
    print(recons_img.shape, gt_img.shape)
    return gt_img, recons_img, all_bit_indices

def load_visual_tokenizer(args):
    device = DEVICE
    # load vae
    if args.vae_type in [16,18,20,24,32,64]:
        from models.bsq_vae.vae import vae_model
        schedule_mode = "dynamic"
        codebook_dim = args.vae_type
        codebook_size = 2**codebook_dim
        if args.apply_spatial_patchify:
            patch_size = 8
            encoder_ch_mult=[1, 2, 4, 4]
            decoder_ch_mult=[1, 2, 4, 4]
        else:
            patch_size = 16
            encoder_ch_mult=[1, 2, 4, 4, 4]
            decoder_ch_mult=[1, 2, 4, 4, 4]
        vae = vae_model(args.vae_path, schedule_mode, codebook_dim, codebook_size, patch_size=patch_size, 
                        encoder_ch_mult=encoder_ch_mult, decoder_ch_mult=decoder_ch_mult, test_mode=True).to(device)
    else:
        raise ValueError(f'vae_type={args.vae_type} not supported')
    return vae

def load_transformer(vae, args):
    model_path = args.model_path
    if args.checkpoint_type == 'torch': 
        # copy large model to local; save slim to local; and copy slim to nas; load local slim model
        if osp.exists(args.cache_dir):
            local_model_path = osp.join(args.cache_dir, 'tmp', model_path.replace('/', '_'))
        else:
            local_model_path = model_path
        if args.enable_model_cache:
            slim_model_path = model_path.replace('ar-', 'slim-')
            local_slim_model_path = local_model_path.replace('ar-', 'slim-')
            os.makedirs(osp.dirname(local_slim_model_path), exist_ok=True)
            print(f'model_path: {model_path}, slim_model_path: {slim_model_path}')
            print(f'local_model_path: {local_model_path}, local_slim_model_path: {local_slim_model_path}')
            if not osp.exists(local_slim_model_path):
                if osp.exists(slim_model_path):
                    print(f'copy {slim_model_path} to {local_slim_model_path}')
                    shutil.copyfile(slim_model_path, local_slim_model_path)
                else:
                    if not osp.exists(local_model_path):
                        print(f'copy {model_path} to {local_model_path}')
                        shutil.copyfile(model_path, local_model_path)
                    save_slim_model(local_model_path, save_file=local_slim_model_path, device=DEVICE)
                    print(f'copy {local_slim_model_path} to {slim_model_path}')
                    if not osp.exists(slim_model_path):
                        shutil.copyfile(local_slim_model_path, slim_model_path)
                        os.remove(local_model_path)
                        os.remove(model_path)
            slim_model_path = local_slim_model_path
        else:
            slim_model_path = model_path
        print(f'load checkpoint from {slim_model_path}')

    model_config = ModelConfig.from_type(args.model_type)
    infinity = load_infinity(
        rope2d_each_sa_layer=args.rope2d_each_sa_layer, 
        rope2d_normalized_by_hw=args.rope2d_normalized_by_hw,
        use_scale_schedule_embedding=args.use_scale_schedule_embedding,
        pn=args.pn,
        use_bit_label=args.use_bit_label, 
        add_lvl_embeding_only_first_block=args.add_lvl_embeding_only_first_block, 
        model_path=slim_model_path, 
        scale_schedule=None, 
        vae=vae, 
        model_kwargs=model_config.to_dict(),
        text_channels=args.text_channels,
        apply_spatial_patchify=args.apply_spatial_patchify,
        use_flex_attn=args.use_flex_attn,
        bf16=args.bf16,
    )
    return infinity

def enhance_prompt(prompt):
    SYSTEM = """You are part of a team of bots that creates images. You work with an assistant bot that will draw anything you say.  

When given a user prompt, your role is to transform it into a creative, detailed, and vivid image description that focuses on visual and sensory features. Avoid directly referencing specific real-world people, places, or cultural knowledge unless explicitly requested by the user.  

### Guidelines for Generating the Output:  

1. **Output Format:**  
   Your response must be in the following dictionary format:  
   ```json
   {
     "prompt": "<enhanced image description>",
     "cfg": <cfg value>
   }
   ```  

2. **Enhancing the "prompt" field:**  
   - Use your creativity to expand short or vague prompts into highly detailed, visually rich descriptions.  
   - Focus on describing visual and sensory elements, such as colors, textures, shapes, lighting, and emotions.  
   - Avoid including known real-world information unless the user explicitly requests it. Instead, describe features that evoke the essence or appearance of the scene or subject.  
   - For particularly long user prompts (over 50 words), output them directly without refinement.  
   - Image descriptions must remain between 8-512 words. Any excess text will be ignored.  
   - If the user's request involves rendering specific text in the image, enclose that text in single quotation marks and prefix it with "the text".  

3. **Determining the "cfg" field:**  
   - If the image to be generated is likely to feature a clear face, set `"cfg": 1`.  
   - If the image does not prominently feature a face, set `"cfg": 3`.  

4. **Examples of Enhanced Prompts:**  
   - **User prompt:** "a tree"  
     **Enhanced prompt:** "A towering tree with a textured bark of intricate ridges and grooves stands under a pale blue sky. Its sprawling branches create an umbrella of rich, deep green foliage, with a few golden leaves scattered, catching the sunlight like tiny stars."  
     **Cfg:** `3`  

   - **User prompt:** "a person reading"  
     **Enhanced prompt:** "A figure sits on a cozy armchair, illuminated by the soft, warm glow of a nearby lamp. Their posture is relaxed, and their hands gently hold an open book. Shadows dance across their thoughtful expression, while the fabric of their clothing appears textured and soft, with subtle folds."  
     **Cfg:** `1`  

5. **Your Output:**  
   Always return a single dictionary containing both `"prompt"` and `"cfg"` fields. Avoid any additional commentary or explanations.  

Don't write anything except the dictionary in the output. (Don't start with ```)
"""
    result = client.predict(
            query=prompt,
            history=[],
            system=SYSTEM,
            api_name="/model_chat"
    )

    dict_of_inputs = json.loads(result[1][-1][-1])
    print(dict_of_inputs)

    return gr.update(value=dict_of_inputs["prompt"]), gr.update(value=float(dict_of_inputs['cfg']))

# Set up paths
weights_path = Path(__file__).parent / 'weights'
weights_path.mkdir(exist_ok=True)
download_infinity_weights(weights_path)

# Define args
args = argparse.Namespace(
    pn='1M',
    model_path=str(weights_path / 'infinity_2b_reg.pth'),
    cfg_insertion_layer=0,
    vae_type=32,
    vae_path=str(weights_path / 'infinity_vae_d32reg.pth'),
    add_lvl_embeding_only_first_block=1,
    use_bit_label=1,
    model_type='infinity_2b',
    rope2d_each_sa_layer=1,
    rope2d_normalized_by_hw=2,
    use_scale_schedule_embedding=0,
    sampling_per_bits=1,
    text_channels=2048,
    apply_spatial_patchify=0,
    h_div_w_template=1.000,
    use_flex_attn=0,
    cache_dir='/dev/shm',
    checkpoint_type='torch',
    seed=0,
    bf16=1 if torch.bfloat16 == torch.get_default_dtype() else 0,
    save_file='tmp.jpg',
    enable_model_cache=False,
)

# Load models
text_tokenizer, text_encoder = load_tokenizer(t5_path="google/flan-t5-xl")
vae = load_visual_tokenizer(args)
infinity = load_transformer(vae, args)

# Define the image generation function
@spaces.GPU
def generate_image(prompt, cfg, tau, h_div_w, seed, enable_positive_prompt=False):
    """Generate an image from a prompt with integrated generation logic."""
    try:
        # Set random seed for reproducibility
        if seed is not None:
            torch.manual_seed(seed)
            random.seed(seed)
            np.random.seed(seed)
        
        # Calculate image dimensions
        tgt_h, tgt_w = dynamic_resolution_h_w(h_div_w)
        scale_schedule = None
        
        # Process text prompt
        text_cond_tuple = encode_prompt(text_tokenizer, text_encoder, prompt, enable_positive_prompt)
        
        # Set up negative prompt if needed
        negative_prompt = ''
        if negative_prompt:
            negative_cond_tuple = encode_prompt(text_tokenizer, text_encoder, negative_prompt)
            negative_label_B_or_BLT = negative_cond_tuple[0]
        else:
            negative_label_B_or_BLT = None
        
        print(f'cfg: {cfg}, tau: {tau}')
        
        # Generate image with automatic mixed precision
        with torch.amp.autocast(device_type=DEVICE.type, dtype=torch.bfloat16):
            stt = time.time()
            _, _, img_list = infinity.autoregressive_infer_cfg(
                vae=vae,
                text_cond_tuple=text_cond_tuple,
                negative_label_B_or_BLT=negative_label_B_or_BLT,
                cfg_list=[cfg],
                tau_list=[tau],
                top_k=900,
                top_p=0.97,
                cfg_sc=3,
                cfg_exp_k=0.0,
                cfg_insertion_layer=[args.cfg_insertion_layer],
                vae_type=args.vae_type,
                gumbel=0,
                softmax_merge_topk=-1,
                gt_leak=0,
                gt_ls_Bl=None,
                g_seed=seed,
                sampling_per_bits=args.sampling_per_bits,
                scale_schedule=scale_schedule,
            )
            print(f'inference time: {time.time()-stt:.3f}s')
        
        # Convert the image efficiently
        with torch.no_grad():
            image = img_list[0].cpu().numpy()
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = np.uint8(image)
        
        return image
    except Exception as e:
        print(f"Error generating image: {e}")
        return None

# Set up Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("<h1><center>Infinity Image Generator</center></h1>")
    gr.Markdown("### Instructions")
    gr.Markdown("1. Enter a prompt in the **Prompt Settings** section.")
    gr.Markdown("2. Click the **Enhance Prompt** button to generate a more creative and detailed prompt.")
    gr.Markdown("3. Adjust the **Image Settings** as desired.")
    gr.Markdown("4. Click the **Generate Image** button to generate the image on the right.")
    
    with gr.Row():
        with gr.Column():
            # Prompt Settings
            gr.Markdown("### Prompt Settings")
            prompt = gr.Textbox(label="Prompt", value="alien spaceship enterprise", placeholder="Enter your prompt here...")
            enhance_prompt_button = gr.Button("Enhance Prompt", variant="secondary")

            # Image Settings
            gr.Markdown("### Image Settings")
            with gr.Row():
                cfg = gr.Slider(label="CFG (Classifier-Free Guidance)", minimum=1, maximum=10, step=0.5, value=3, info="Controls the strength of the prompt.")
                tau = gr.Slider(label="Tau (Temperature)", minimum=0.1, maximum=1.0, step=0.1, value=0.5, info="Controls the randomness of the output.")
            with gr.Row():
                h_div_w = gr.Slider(label="Aspect Ratio (Height/Width)", minimum=0.5, maximum=2.0, step=0.1, value=1.0, info="Set the aspect ratio of the generated image.")
                seed = gr.Number(label="Seed", value=random.randint(0, 10000), info="Set a seed for reproducibility.")
            
            # Generate Button
            generate_button = gr.Button("Generate Image", variant="primary")
        
        with gr.Column():
            # Output Section
            gr.Markdown("### Generated Image")
            output_image = gr.Image(label="Generated Image", type="pil")
    
    # Error Handling
    error_message = gr.Textbox(label="Error Message", visible=False)
    
    # Link the enhance prompt button to the prompt enhancement function
    enhance_prompt_button.click(
        enhance_prompt,
        inputs=prompt,
        outputs=[prompt, cfg],
    )

    # Link the generate button to the image generation function
    generate_button.click(
        generate_image,
        inputs=[prompt, cfg, tau, h_div_w, seed],
        outputs=output_image
    )

# Launch the Gradio app
demo.launch()