allenai / allenai/unified-io-2.pytorch

Simple test_demo error using unified-io2 for text to image generation

オープン
#14 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
Python
スター
78
フォーク
8
PR マージ指標
30日以内にマージされた PR はありません

説明

Hi, all. I used the code below to test a small example of text to image generation using uni-io2. Where could possibly the error come from? The output image seems like it is not generating a plausible output.

```
import torch
import numpy as np
from pathlib import Path
from PIL import Image
from uio2.model import UnifiedIOModel
from uio2.preprocessing import UnifiedIOPreprocessor, build_batch
from transformers import LlamaTokenizer

def main():

prompt = "Generate an image of a bunny"

output_dir = Path("results/unified_io2_test")
output_dir.mkdir(parents=True, exist_ok=True)

model = UnifiedIOModel.from_pretrained("allenai/uio2-large")
tokenizer = LlamaTokenizer.from_pretrained("./tokenizer.model")
preprocessor = UnifiedIOPreprocessor.from_pretrained("allenai/uio2-preprocessor", tokenizer=tokenizer)

model.to("cuda" if torch.cuda.is_available() else "cpu")
model.set_modalities(input_modalities=["text"], target_modalities=["image"])

try:
preprocessed = preprocessor(text_inputs=prompt, target_modality="image")
with torch.no_grad():
batch = build_batch([preprocessed], device=model.device)
result = model.generate(batch, modality="image")
if hasattr(result, 'detach'):
result = result.detach().cpu().numpy()
result = np.squeeze(result)
if result.dtype != np.uint8:
result = np.clip(result * 255, 0, 255).astype(np.uint8) if result.max() <= 1.0 else result.astype(np.uint8)
img = Image.fromarray(result)
img.save(output_dir / "test_output.png")
print(f"[Success] Saved generated image to {output_dir / 'test_output.png'}")
except Exception as e:
print(f"[Error] Failed to generate image: {e}")

if __name__ == "__main__":
main()
```

Image

コントリビューションガイド

このリポジトリのコントリビューションガイドは索引されていません

評価

この issue はまだ評価されていません。

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。