app.py

import spaces
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from huggingface_hub import login
import os

# Загрузка модели с использованием ZeroGPU
@spaces.GPU
def load_model():
    model_id = "CultriX/flux-nsfw-highress"
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        use_safetensors=True,  # Явное указание использовать SafeTensors
        torch_dtype=torch.float16  # Использование половинной точности для экономии памяти
    )
    pipe = pipe.to("cuda")
    return pipe

# Функция для генерации изображения
@spaces.GPU(duration=120)  # Указываем длительность до 120 секунд
def generate_image(prompt, negative_prompt="", guidance_scale=7.5, steps=30):
    pipe = load_model()
    
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=steps
    ).images[0]
    
    return image

# Создание интерфейса Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Демонстрация модели flux-nsfw-highress")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Запрос", placeholder="Введите описание желаемого изображения")
            negative_prompt = gr.Textbox(label="Негативный запрос", placeholder="Что исключить из изображения")
            guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.5, label="Guidance Scale")
            steps = gr.Slider(minimum=10, maximum=50, value=30, step=1, label="Количество шагов")
            generate_btn = gr.Button("Сгенерировать")
        
        with gr.Column():
            output_image = gr.Image(label="Сгенерированное изображение")
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, guidance_scale, steps],
        outputs=output_image
    )

# Запуск приложения
demo.launch()