1import torch
2from diffusers import AutoPipelineForText2Image
3from transformers import CLIPTokenizer, CLIPTextModel
4
5device = "cuda" if torch.cuda.is_available() else "cpu"
6lang = "hin_Deva" # Hindi
7
8# Load pipeline
9pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
10
11# Load the multilingual tokenizer
12tokenizer = CLIPTokenizer.from_pretrained("tokenizers/multilingual")
13pipe.tokenizer = tokenizer
14pipe.text_encoder.resize_token_embeddings(len(tokenizer))
15
16# Load the fine-tuned text encoder
17state_dict = torch.load(f"models/{lang}/{lang}_text_encoder.pth")
18new_text_encoder = CLIPTextModel(config=pipe.text_encoder.config)
19new_text_encoder.load_state_dict(state_dict)
20new_text_encoder = new_text_encoder.to(device)
21pipe.text_encoder = new_text_encoder
22pipe = pipe.to(device)
23
24# Generate and save image
25caption = "गाँव का शांतिपूर्ण दृश्य|"
26image = pipe(caption).images[0]
27image.save(f"example.png")
