Spaces:
Sleeping
Sleeping
File size: 1,063 Bytes
870ab35 4c63345 870ab35 4c63345 870ab35 4c63345 870ab35 4c63345 870ab35 4c63345 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
# app.py
import io
import base64
from fastapi import FastAPI
from pydantic import BaseModel
from diffusers import StableDiffusionPipeline
from PIL import Image
import torch
app = FastAPI()
# Load model on CPU
pipe = StableDiffusionPipeline.from_pretrained(
"SG161222/Realistic_Vision_V5.1_noVAE",
torch_dtype=torch.float32, # use float32 for CPU
safety_checker=None,
use_safetensors=True
).to("cpu") # switched from "cuda" to "cpu"
class PromptRequest(BaseModel):
prompt: str
negative_prompt: str = ""
width: int = 512
height: int = 512
@app.post("/generate")
def generate_image(data: PromptRequest):
image = pipe(
prompt=data.prompt,
negative_prompt=data.negative_prompt,
width=data.width,
height=data.height
).images[0]
buffered = io.BytesIO()
image.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
return {"image_base64": img_str}
@app.get("/")
def root():
return {"message": "Image generation API running on CPU!"}
|