import gradio as gr
import torch
import os
import secrets
from diffusers import StableDiffusionPipeline, LCMScheduler, AutoencoderTiny
from huggingface_hub import hf_hub_download

# Папка для файлов
os.makedirs("outputs", exist_ok=True)

# 1. Загрузка модели
model_path = hf_hub_download(repo_id="zw89/DreamShaper", filename="DreamShaper_8_pruned.safetensors")
pipe = StableDiffusionPipeline.from_single_file(model_path, torch_dtype=torch.float32)

# 2. Ускорители
lora_path = hf_hub_download(repo_id="zw89/lcm-lora-sdv1-5", filename="pytorch_lora_weights.safetensors")
pipe.load_lora_weights(lora_path)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to("cpu")

# 1. Изменяем функцию, чтобы она принимала ДВА аргумента
def generate(prompt, negative_prompt):
    seed = secrets.randbits(32)
    generator = torch.Generator().manual_seed(seed)
    
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=5,
        guidance_scale=1.0,
        width=512, 
        height=512,
        generator=generator
    ).images[0]
    
    file_path = f"outputs/art_{seed}.png"
    image.save(file_path, "PNG")
    return file_path

# 2. Настраиваем интерфейс с двумя полями
demo = gr.Interface(
    fn=generate, 
    inputs=[
        gr.Textbox(label="Promt"),
        gr.Textbox(label="Negative", value="low quality, plastic, blurry") # Значение по умолчанию
    ],
    outputs=gr.Image(type="filepath", label="PNG"),
    title="SD Mini Dream"
)


demo.launch(server_name="0.0.0.0", server_port=7860)