import gradio as gr
from diffusers import StableDiffusionPipeline

model_id = "cagliostrolab/animagine-xl-3.1"
pipe = StableDiffusionPipeline.from_pretrained(model_id)

def generate_image(prompt, negative_prompt, seed, cfg_scale, strength, steps):
  if seed == -1:
    seed = None
  generator = torch.manual_seed(seed)
  image = pipe(
      prompt,
      negative_prompt=negative_prompt,
      guidance_scale=cfg_scale,
      num_inference_steps=int(steps),
      strength=strength,
      generator=generator,
  ).images[0]
  return image

with gr.Blocks() as demo:
  with gr.Tabs():
    with gr.Tab("Основные настройки"):
      with gr.Row():
        prompt = gr.Textbox(label="Подсказка")
    with gr.Tab("Расширенные настройки"):
      with gr.Row():
        negative_prompt = gr.Textbox(label="Негативная подсказка")
        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=-1)
        cfg_scale = gr.Slider(label="CFG", minimum=1.0, maximum=15.0, step=0.5, value=7.5)
        strength = gr.Slider(label="Strength", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
        steps = gr.Slider(label="Sampling steps", minimum=1, maximum=100, step=1, value=50)

  btn = gr.Button("Сгенерировать")
  output = gr.Image(label="Результат")

  btn.click(fn=generate_image, inputs=[prompt, negative_prompt, seed, cfg_scale, strength, steps], outputs=output)

demo.launch()