routers/training.py

from fastapi import APIRouter, Form, BackgroundTasks
from config import settings
import os
import json
from routers.donut_evaluate import run_evaluate_donut
from routers.donut_training import run_training_donut
import utils
import torch
from diffusers import StableDiffusionPipeline

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("radames/stable-diffusion-v1-5-img2img")
# pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to("cuda")


router = APIRouter()

@router.get("/hi")
async def hifunction():
    prompt = "a photograph of an astronaut riding a horse"
    image = pipe(prompt).images[0]
    print(image)
    return ["done"]