allenai / allenai/unified-io-2.pytorch
Simple test_demo error using unified-io2 for text to image generation
- 主要言語
- 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()
```
コントリビューションガイド
このリポジトリのコントリビューションガイドは索引されていません
評価
この issue はまだ評価されていません。