import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Выбираем модель, которая точно поддерживает CPU
model_id = "runwayml/stable-diffusion-v1-5" 

# Отключаем использование высокопроизводительных инструкций xFormers, так как они не нужны на CPU
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_auth_token=True, variant="fp16", revision="fp16")

# Не перемещаем модель на GPU
# pipe = pipe.to("cuda") 

def generate_image(prompt, guidance_scale=7.5, num_inference_steps=50):
  image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images[0]
  return image

with gr.Blocks() as demo:
  with gr.Tab("Генерация"):
    with gr.Row():
      with gr.Column():
        prompt_input = gr.Textbox(label="Введите описание картинки", lines=2)
        generate_button = gr.Button(label="Сгенерировать")
      with gr.Column():
        image_output = gr.Image(label="Результат")
  with gr.Tab("Настройки"):
    with gr.Row():
      guidance_scale_slider = gr.Slider(label="Guidance Scale (чем больше, тем больше соответствует описанию)", minimum=1, maximum=20, step=0.5, value=7.5)
      num_inference_steps_slider = gr.Slider(label="Number of Inference Steps (чем больше, тем качественнее, но дольше)", minimum=10, maximum=100, step=10, value=50)

  generate_button.click(fn=generate_image, 
                       inputs=[prompt_input, guidance_scale_slider, num_inference_steps_slider], 
                       outputs=image_output)

demo.launch()