import gradio as gr
from gradio import StableDiffusion, DPMSolverMultistepScheduler
import torch
from transformers import AutoProcessor, AutoModelForCausalLM

# Загрузка моделей
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe_sd = StableDiffusion.from_pretrained("CompVis/ldm-text2im-large-256").to(device)
scheduler = DPMSolverMultistepScheduler.from_config(pipe_sd.scheduler.config)
pipe_sd.scheduler = scheduler

processor_dalle = AutoProcessor.from_pretrained("openai/dall-e-3")
model_dalle = AutoModelForCausalLM.from_pretrained("openai/dall-e-3", device_map="auto")

def generate_with_stable_diffusion(prompt, num_inference_steps, guidance_scale, seed):
    generator = torch.Generator(device).manual_seed(seed)
    image = pipe_sd([prompt], num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator).images[0]
    return image

def generate_with_dalle(prompt, num_images):
    input_ids = processor_dalle.tokenize(prompt, padding="MAX_LENGTH", max_length=128, truncation=True, return_tensors="pt").input_ids
    output = model_dalle.generate(input_ids.to(model_dalle.device), max_new_tokens=256, do_sample=True, top_p=0.95, temperature=1.0, num_return_sequences=num_images)
    images = processor_dalle.batch_decode(output, skip_special_tokens=True)
    return images

def main():
    with gr.Blocks() as demo:
        gr.Markdown("# Генератор изображений с использованием Stable Diffusion и DALL-E")
        with gr.Tab("Stable Diffusion"):
            with gr.Row():
                prompt_input = gr.Textbox(label="Описание изображения", placeholder="Введите текстовое описание изображения...")
                num_inference_steps_slider = gr.Slider(minimum=1, maximum=50, step=1, label="Количество шагов вывода", value=25)
                guidance_scale_slider = gr.Slider(minimum=1, maximum=30, step=1, label="Масштаб руководства", value=7.5)
                seed_input = gr.Number(label="Значение случайной генерации (seed)", value=42)
                
            with gr.Row():
                generate_button = gr.Button("Создать изображение")
                image_output = gr.Image(label="Генерируемое изображение")

            generate_button.click(generate_with_stable_diffusion, inputs=[prompt_input, num_inference_steps_slider, guidance_scale_slider, seed_input], outputs=image_output)

        with gr.Tab("DALL-E"):
            with gr.Row():
                prompt_input_dalle = gr.Textbox(label="Описание изображения", placeholder="Введите текстовое описание изображения...")
                num_images_slider = gr.Slider(minimum=1, maximum=4, step=1, label="Количество генерируемых изображений", value=1)

            with gr.Row():
                generate_button_dalle = gr.Button("Создать изображение")
                image_output_dalle = gr.Gallery(label="Генерируемые изображения")

            generate_button_dalle.click(generate_with_dalle, inputs=[prompt_input_dalle, num_images_slider], outputs=image_output_dalle)

    demo.launch()

if __name__ == "__main__":
    main()