from fastapi import FastAPI
from pydantic import BaseModel
import requests  # Добавляем импорт
import asyncio
from contextlib import asynccontextmanager
from diffusers import DiffusionPipeline

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Здесь можно инициализировать ресурсы
    yield
    # Здесь можно освободить ресурсы
    
app = FastAPI(lifespan=lifespan)



class ImageRequest(BaseModel):  # Новая модель для запроса
    prompt: str = "girl in red dress"

@app.post("/generate-image")
async def generate_image(request: ImageRequest):    
    try:
        pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell")
        
        prompt = request.prompt
        image = pipe(prompt).images[0]        
        
        return image
    except Exception as e:
        return {"error": str(e)}