import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

class TextToImageGenerator(torch.nn.Module):
    def __init__(self, model_name="gpt2"):
        super(TextToImageGenerator, self).__init__()
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        self.gpt2 = GPT2LMHeadModel.from_pretrained(model_name)

    def forward(self, input_text):
        input_ids = self.tokenizer(input_text, return_tensors="pt")["input_ids"]
        output = self.gpt2(input_ids, return_dict=True)
        return output.logits

# Instanciar el modelo
model = TextToImageGenerator()

# Imprimir la arquitectura del modelo
print(model)
input_text = "Una escena de montaña nevada al atardecer"
image_logits = model(input_text)

# Aquí, deberías tener una capa de salida que represente tu imagen generada.
# Podrías utilizar técnicas más avanzadas y adaptar este ejemplo según tus necesidades.