Could you prepare a reasoning demo?

[win10ogod]

Could you prepare a reasoning demo?
I have some doubts about the reasoning.
I would like to ask you to help me check if there are any errors in training and inference.
train.py

import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from transformers import AutoTokenizer, get_linear_schedule_with_warmup
from transfusion_pytorch import Transfusion
from bitsandbytes.optim import AdamW8bit
from torch.amp import autocast, GradScaler
from tqdm import tqdm

class ImageTextDataset(Dataset):
    def __init__(self, data_dir, tokenizer, max_length=64, image_size=256):
        self.data_dir = data_dir
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.image_transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
        ])
        self.samples = [f for f in os.listdir(data_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        image_file = self.samples[idx]
        text_file = os.path.splitext(image_file)[0] + '.txt'
        
        image_path = os.path.join(self.data_dir, image_file)
        text_path = os.path.join(self.data_dir, text_file)
        
        image = Image.open(image_path).convert('RGB')
        image = self.image_transform(image)
        
        with open(text_path, 'r', encoding='utf-8') as f:
            text = f.read().strip()
        
        encoded_text = self.tokenizer.encode(text, max_length=self.max_length, padding='max_length', truncation=True, return_tensors='pt').squeeze(0)
        
        return encoded_text, image

def create_model(vocab_size, image_size):
    model = Transfusion(
        num_text_tokens=vocab_size,
        dim_latent=384,
        channel_first_latent=True,
        modality_default_shape=(image_size // 8, image_size // 8),
        modality_encoder=torch.nn.Sequential(
            torch.nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
        ),
        modality_decoder=torch.nn.Sequential(
            torch.nn.ConvTranspose2d(384, 256, kernel_size=3, stride=2, padding=1, output_padding=1),
            torch.nn.ReLU(),
            torch.nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
            torch.nn.ReLU(),
            torch.nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1),
            torch.nn.Sigmoid(),
        ),
        transformer=dict(
            dim=512,
            depth=8
        )
    )
    # 確保所有參數都設置了 requires_grad=True
    for param in model.parameters():
        param.requires_grad = True
    return model

def train_epoch(model, train_loader, optimizer, scheduler, scaler, device, gradient_accumulation_steps, max_grad_norm):
    model.train()
    total_loss = 0
    pbar = tqdm(train_loader, desc="Training")
    
    for batch_idx, (texts, images) in enumerate(pbar):
        texts, images = texts.to(device), images.to(device)
        
        if batch_idx % gradient_accumulation_steps == 0:
            optimizer.zero_grad()
        
        with autocast(device_type='cuda', dtype=torch.float16):
            batch = [[text, image] for text, image in zip(texts, images)]
            try:
                # 直接調用模型，不使用梯度檢查點
                loss = model(batch)
            except Exception as e:
                print(f"Error in model forward pass: {e}")
                continue
        
        if not torch.isfinite(loss):
            print(f"Warning: Loss is {loss}, skipping batch")
            continue
        
        scaled_loss = scaler.scale(loss / gradient_accumulation_steps)
        
        try:
            scaled_loss.backward()
        except RuntimeError as e:
            print(f"Error in backward pass: {e}")
            print("Model parameters requiring grad:")
            for name, param in model.named_parameters():
                print(f"{name}: requires_grad = {param.requires_grad}")
            continue
        
        if (batch_idx + 1) % gradient_accumulation_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
        
        total_loss += loss.item()
        pbar.set_postfix({'loss': loss.item()})
    
    return total_loss / len(train_loader)

def train(model, train_loader, optimizer, scheduler, scaler, device, epochs, gradient_accumulation_steps, max_grad_norm, checkpoint_dir):
    for epoch in range(epochs):
        avg_loss = train_epoch(model, train_loader, optimizer, scheduler, scaler, device, gradient_accumulation_steps, max_grad_norm)
        print(f'Epoch {epoch+1}/{epochs}, Average Loss: {avg_loss:.4f}')
        
        # 保存檢查點
        checkpoint_path = os.path.join(checkpoint_dir, f'checkpoint_epoch_{epoch+1}.pt')
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
            'loss': avg_loss,
        }, checkpoint_path)
        print(f'Checkpoint saved to {checkpoint_path}')

def main():
    data_dir = 'E:/data/afg/1_zkz'
    image_size = 256
    batch_size = 4
    learning_rate = 1e-4
    epochs = 30
    warmup_steps = 1000
    gradient_accumulation_steps = 4
    max_grad_norm = 1.0
    checkpoint_dir = 'checkpoints'
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    os.makedirs(checkpoint_dir, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-0.5B-Instruct')
    
    dataset = ImageTextDataset(data_dir, tokenizer, image_size=image_size)
    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    
    model = create_model(vocab_size=tokenizer.vocab_size, image_size=image_size)
    model = model.to(device)
    
    optimizer = AdamW8bit(model.parameters(), lr=learning_rate)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=len(train_loader) * epochs)
    scaler = GradScaler()

    # 檢查是否有現有的檢查點
    latest_checkpoint = max([f for f in os.listdir(checkpoint_dir) if f.startswith('checkpoint')], default=None)
    if latest_checkpoint:
        checkpoint_path = os.path.join(checkpoint_dir, latest_checkpoint)
        checkpoint = torch.load(checkpoint_path)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        start_epoch = checkpoint['epoch'] + 1
        print(f'Resuming from checkpoint {checkpoint_path}')
    else:
        start_epoch = 0

    train(model, train_loader, optimizer, scheduler, scaler, device, epochs - start_epoch, gradient_accumulation_steps, max_grad_norm, checkpoint_dir)

if __name__ == '__main__':
    main()

inference.py

import os
import torch
import gradio as gr
from PIL import Image
from torchvision import transforms
from transformers import AutoTokenizer
from transfusion_pytorch import Transfusion, print_modality_sample
from torch.amp import autocast

def create_model(vocab_size, image_size):
    model = Transfusion(
        num_text_tokens=vocab_size,
        dim_latent=384,
        channel_first_latent=True,
        modality_default_shape=(image_size // 8, image_size // 8),
        modality_encoder=torch.nn.Sequential(
            torch.nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
        ),
        modality_decoder=torch.nn.Sequential(
            torch.nn.ConvTranspose2d(384, 256, kernel_size=3, stride=2, padding=1, output_padding=1),
            torch.nn.ReLU(),
            torch.nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
            torch.nn.ReLU(),
            torch.nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1),
            torch.nn.Sigmoid(),
        ),
        transformer=dict(
            dim=512,
            depth=8
        )
    )
    return model

# 初始化模型和分词器
tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-0.5B-Instruct')
model = create_model(vocab_size=tokenizer.vocab_size, image_size=256)

# 加载训练好的模型
checkpoint_dir = 'checkpoints'
latest_checkpoint = max([f for f in os.listdir(checkpoint_dir) if f.startswith('checkpoint')], default=None)
if latest_checkpoint:
    checkpoint_path = os.path.join(checkpoint_dir, latest_checkpoint)
    checkpoint = torch.load(checkpoint_path, map_location='gpu')
    model.load_state_dict(checkpoint['model_state_dict'])
    print(f'Loaded checkpoint: {checkpoint_path}')
else:
    print('No checkpoint found. Using untrained model.')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()

# 图像预处理
image_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

def generate(text_input, image_input=None, max_length=512):
    with torch.no_grad(), autocast(device_type='cuda', dtype=torch.float16):
        try:
            # Encode text or process image input
            if text_input:
                encoded_text = tokenizer.encode(text_input, max_length=max_length, padding='max_length', truncation=True, return_tensors='pt').squeeze(0).to(device)
                sample = model.sample(encoded_text, max_length=max_length)
            elif image_input is not None:
                image_tensor = image_transform(image_input).unsqueeze(0).to(device)
                sample = model.sample(image_tensor)
            else:
                sample = model.sample()

            # Check sample output
            if sample is None:
                print("Warning: No sample generated.")
                return "No output generated", None

            print("Sample raw output:", sample)
            
            generated_text, generated_image = None, None

            # Extract and handle modalities from the sample
            for item in sample:
                if isinstance(item, tuple):
                    modality_type, modality_data = item
                    if modality_type == 0:  # Assuming modality_type 0 means image
                        print(f"Image modality data found with shape: {modality_data.shape}")
                        generated_image = modality_data
                elif isinstance(item, torch.Tensor) and item.dtype in (torch.int, torch.long):
                    print(f"Text modality data found with shape: {item.shape}")
                    generated_text = item.tolist()

            # Decode text output if available
            if generated_text is not None:
                generated_text = tokenizer.decode(generated_text)
                print("Generated text:", generated_text)

            # Convert the image tensor to a PIL image if available
            if generated_image is not None:
                if generated_image.dim() == 3:
                    generated_image = generated_image.squeeze().permute(1, 2, 0).cpu().numpy()
                    generated_image = Image.fromarray((generated_image * 255).astype('uint8'))
                else:
                    print("Unexpected image data dimensions:", generated_image.shape)
                    generated_image = None

            return generated_text, generated_image
        
        except Exception as e:
            print("Error during generation:", str(e))
            return f"Error: {str(e)}", None

# 单独的功能函数
def text_to_image(text_input):
    generated_text, generated_image = generate(text_input)
    return generated_image, generated_text

def image_to_image(image_input):
    generated_text, generated_image = generate(None, image_input)
    return generated_image, generated_text

def text_to_text(text_input):
    generated_text, generated_image = generate(text_input)
    return generated_text, generated_image

def text_and_image_to_both(text_input, image_input):
    generated_text, generated_image = generate(text_input, image_input)
    return generated_text, generated_image

# 使用 Blocks 创建界面
with gr.Blocks(title="Transfusion: 多模态生成演示") as demo:
    gr.Markdown("# Transfusion: 多模态生成演示")
    gr.Markdown("支持文生图、图生图、文本交互生成和图文混合生成。")

    with gr.Tab("文本到图像"):
        text_input1 = gr.Textbox(lines=3, label="输入文本")
        text_to_image_btn = gr.Button("生成图像")
        image_output1 = gr.Image(type="pil", label="生成的图像")
        text_output1 = gr.Textbox(label="生成的文本描述")
        text_to_image_btn.click(text_to_image, inputs=text_input1, outputs=[image_output1, text_output1])

    with gr.Tab("图像到图像"):
        image_input = gr.Image(type="pil", label="输入图像")
        image_to_image_btn = gr.Button("生成新图像")
        image_output2 = gr.Image(type="pil", label="生成的图像")
        text_output2 = gr.Textbox(label="生成的文本描述")
        image_to_image_btn.click(image_to_image, inputs=image_input, outputs=[image_output2, text_output2])

    with gr.Tab("文本到文本"):
        text_input2 = gr.Textbox(lines=3, label="输入文本")
        text_to_text_btn = gr.Button("生成文本")
        text_output3 = gr.Textbox(label="生成的文本")
        image_output3 = gr.Image(type="pil", label="生成的相关图像")
        text_to_text_btn.click(text_to_text, inputs=text_input2, outputs=[text_output3, image_output3])

    with gr.Tab("文本和图像到两者"):
        text_input3 = gr.Textbox(lines=3, label="输入文本")
        image_input2 = gr.Image(type="pil", label="输入图像")
        both_to_both_btn = gr.Button("生成文本和图像")
        text_output4 = gr.Textbox(label="生成的文本")
        image_output4 = gr.Image(type="pil", label="生成的图像")
        both_to_both_btn.click(text_and_image_to_both, inputs=[text_input3, image_input2], outputs=[text_output4, image_output4])

demo.launch()

[lavinal712]

Did you train the model? Does it work?

[win10ogod]

It seems to be effective, but it lacks pure text training. It would be better if I improve the logic of my training script. The current training script is full of errors.

[lavinal712]

Thanks for your answer! I also have some questions about Transfusion.

1. Can the text ability in LLM like LLaMA be transferred into Transfusion? If the model doesn't have text ability, it seems like a diffusion transformer model with bi-attention.
2. Can the model process batched data? I find the code can't load data in batches.
3. Do you think the model architecture like Transfusion, Show-o, Omnigen, can be the new paradigm of image generation and understanding? Maybe it is better to let LLM and Diffusion models do their own jobs.

[win10ogod]

Thanks for your answer! I also have some questions about Transfusion.

1. Can the text ability in LLM like LLaMA be transferred into Transfusion? If the model doesn't have text ability, it seems like a diffusion transformer model with bi-attention.
2. Can the model process batched data? I find the code can't load data in batches.
3. Do you think the model architecture like Transfusion, Show-o, Omnigen, can be the new paradigm of image generation and understanding? Maybe it is better to let LLM and Diffusion models do their own jobs.

The code does not support most common features (yes, including batch data), and model architectures like Transfusion, Show-o, and Omnigen can become a new paradigm for image generation and understanding, but they are not the best answer.
Also, if someone can implement a trainer based on this repository, I'm sure it will be much better than my script, because as a quick test script, I wrote all the code using Claude 3.5s.

[taeinkwon]

Does this code generate images? I tried this code, but it only returns text without an image.

[win10ogod]

Does this code generate images? I tried this code, but it only returns text without an image.

cannot be generated