import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# Загрузка предобученной модели Stable Diffusion
model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda" if torch.cuda.is_available() else "cpu"

try:
    with torch.no_grad():
        pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)
        pipe = pipe.to(device)
except Exception as e:
    print(f"Error loading the model: {e}")
    raise

# Функция для генерации изображения
def generate_image(prompt: str, num_inference_steps: int = 50, guidance_scale: float = 7.5):
    """
    Generate an image using Stable Diffusion.

    Args:
        prompt (str): The text prompt for the image generation.
        num_inference_steps (int): The number of inference steps. Default is 50.
        guidance_scale (float): The guidance scale for the generation. Default is 7.5.
    """
    if not prompt:
        raise ValueError("Prompt cannot be empty")
    if num_inference_steps < 1 or num_inference_steps > 100:
        raise ValueError("Number of inference steps must be between 1 and 100")
    if guidance_scale < 1.0 or guidance_scale > 15.0:
        raise ValueError("Guidance scale must be between 1.0 and 15.0")

    try:
        with torch.no_grad():
            image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
    except Exception as e:
        print(f"Error generating image: {e}")
        raise

    return image

# Создание интерфейса Gradio
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),
        gr.Slider(minimum=1, maximum=100, step=1, value=50, label="Number of Inference Steps"),
        gr.Slider(minimum=1.0, maximum=15.0, step=0.1, value=7.5, label="Guidance Scale")
    ],
    outputs=gr.Image(type="pil"),
    title="Stable Diffusion Image Generator",
    description="Generate images using the Stable Diffusion model."
)

# Запуск интерфейса
try:
    iface.launch(share=True)
except Exception as e:
    print(f"Error launching the interface: {e}")
    raise