Spaces:
Runtime error
Runtime error
File size: 2,375 Bytes
33d3090 0251672 33d3090 0251672 33d3090 7d398da 33d3090 0251672 7d398da 0251672 33d3090 0251672 33d3090 0251672 33d3090 0251672 33d3090 0251672 33d3090 d2b9f1c |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# app.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from diffusers import AutoPipelineForText2Image
from contextlib import asynccontextmanager
import io
import base64
import os
# --- Pydantic Models ---
class ImageRequest(BaseModel):
prompt: str
negative_prompt: str = ""
steps: int = 20 # Lowered for faster CPU inference
class ImageResponse(BaseModel):
image_base64: str
# --- App State and Lifespan ---
app_state = {}
@asynccontextmanager
async def lifespan(app: FastAPI):
# Load the model on startup
hf_token = os.getenv("HF_TOKEN") # Still needed to accept terms
if not hf_token:
raise RuntimeError("HF_TOKEN environment variable not set! Please add it in the Space settings.")
# --- *** THESE ARE THE CHANGES FOR CPU *** ---
# 1. Use the smaller Stable Diffusion v1.5 model
model_id = "runwayml/stable-diffusion-v1-5"
print(f"Loading model: {model_id} for CPU...")
# 2. Load the pipeline without GPU-specific settings
pipe = AutoPipelineForText2Image.from_pretrained(
model_id,
token=hf_token
)
# Note: We do not use .to("cuda")
# --- ************************************** ---
app_state["pipe"] = pipe
print("Model loaded successfully onto CPU.")
yield
# Clean up on shutdown
app_state.clear()
print("Resources cleaned up.")
# --- FastAPI App ---
app = FastAPI(lifespan=lifespan)
@app.get("/")
def root():
return {"status": "Text-to-Image CPU API is running"}
@app.post("/generate-image", response_model=ImageResponse)
def generate_image(request: ImageRequest):
if "pipe" not in app_state:
raise HTTPException(status_code=503, detail="Model is not ready.")
pipe = app_state["pipe"]
print(f"Generating image for prompt: '{request.prompt}'")
try:
image = pipe(
prompt=request.prompt,
negative_prompt=request.negative_prompt,
num_inference_steps=request.steps
).images[0]
buffer = io.BytesIO()
image.save(buffer, format="PNG")
img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
return ImageResponse(image_base64=img_str)
except Exception as e:
print(f"Error during image generation: {e}")
raise HTTPException(status_code=500, detail=str(e)) |