app.py

import torch
from diffusers import AutoPipelineForText2Image
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from transformers import CLIPImageProcessor
import gradio as gr
import numpy as np
import os

# 1. Настройка модели
model_id = "Abobasnik/HiperAi-ImageGenerator"

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

print(f"⏳ Загрузка модели {model_id} на {device}...")

# Загружаем модель
pipe = AutoPipelineForText2Image.from_pretrained(
    model_id, 
    torch_dtype=dtype, 
    variant="fp16" if torch.cuda.is_available() else None
).to(device)

# Загружаем фильтр цензуры
safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to(device)
feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")

def generate_safe(prompt, steps, guidance):
    if not prompt:
        return None, "Введите запрос"
    
    # Генерация с использованием слайдеров
    image = pipe(
        prompt=prompt,
        negative_prompt="nude, nsfw, naked, breasts, gore, violence",
        num_inference_steps=int(steps),
        guidance_scale=float(guidance)
    ).images[0]
    
    # Проверка на цензуру
    img_array = np.array(image)
    clip_input = feature_extractor(images=img_array, return_tensors="pt").to(device)
    _, has_nsfw_concept = safety_checker(images=clip_input.pixel_values, clip_input=clip_input.pixel_values)
    
    if has_nsfw_concept[0]:
        black_img = np.zeros((512, 512, 3), dtype=np.uint8)
        return black_img, "🚨 КОНТЕНТ ЗАБЛОКИРОВАН"
    
    return image, "✅ Готово"

# 2. Интерфейс Gradio
with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
    gr.HTML("<h1 style='text-align: center;'>🛡️ HiperAi-ImageGenerator v1</h1>")
    gr.HTML("<p style='text-align: center;'>Создатель: <b>Graznaa_Lava</b> | Модель: <b>HiperAi-ImageGenerator</b></p>")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Ваш запрос", 
                placeholder="Например: Cute cat in space...",
                lines=3
            )
            # Добавленные настройки
            steps_slider = gr.Slider(1, 12, value=8, step=1, label="Повторы (Steps)")
            guidance_slider = gr.Slider(0.0, 5.0, value=1.5, step=0.1, label="Строгость (Guidance)")
            
            gen_btn = gr.Button("ГЕНЕРИРОВАТЬ", variant="primary")
        
        with gr.Column():
            output_img = gr.Image(label="Результат")
            status_text = gr.Textbox(label="Статус", interactive=False)

    # Обновленные входы (inputs) для генерации
    gen_btn.click(
        fn=generate_safe, 
        inputs=[input_text, steps_slider, guidance_slider], 
        outputs=[output_img, status_text]
    )

demo.launch()