import gradio as gr
from transformers import AutoTokenizer, AutoModelWithLMHead

# Замените 'model_name' на имя вашей модели
model_name = 'facebook/ic_gan'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelWithLMHead.from_pretrained(model_name)

def generate_image(prompt):
    inputs = tokenizer.encode(prompt, return_tensors='pt')
    outputs = model.generate(inputs, max_length=150, temperature=0.7, num_return_sequences=1)
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return text

iface = gr.Interface(fn=generate_image, 
                     inputs=[gr.inputs.Radio(['IC-GAN'], label='Выберите модель'), gr.inputs.Textbox(lines=5, label='Введите промпт')],
                     outputs='text')

iface.launch()