Spaces:

mtzdmir
/

ImageGen

Sleeping

SyntheticVis-Umut

Add specific parameter handling for SD Turbo model

25042b6 about 1 year ago

2.26 kB

import io, os
from typing import Optional
from fastapi import FastAPI
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from huggingface_hub import InferenceClient

app = FastAPI(title="ImageGen API")

# Choose a model that supports text->image via Inference API.
# Using a model that works well with InferenceClient
MODEL_ID = os.getenv("MODEL", "stabilityai/stable-diffusion-xl-base-1.0")
HF_TOKEN = os.getenv("HF_TOKEN")  # set in Space Secrets

# Initialize client with or without token
if HF_TOKEN:
    client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
else:
    # Use public inference API (may have rate limits)
    client = InferenceClient(model=MODEL_ID)

class GenReq(BaseModel):
    prompt: str
    width: int = 768
    height: int = 768
    seed: Optional[int] = None
    num_inference_steps: Optional[int] = None  # some models ignore this

@app.get("/")
def root():
    return {"ok": True, "message": "API is live", "model": MODEL_ID}

@app.post("/generate")
def generate(req: GenReq):
    try:
        # Different models have different parameter requirements
        if "sd-turbo" in MODEL_ID.lower():
            # SD Turbo works with minimal parameters
            img = client.text_to_image(
                prompt=req.prompt,
                num_inference_steps=1,  # SD Turbo works best with 1 step
                guidance_scale=0.0,  # SD Turbo doesn't use guidance scale
                width=req.width,
                height=req.height
            )
        else:
            # Standard SDXL parameters
            img = client.text_to_image(
                prompt=req.prompt,
                width=req.width,
                height=req.height,
                seed=req.seed,
                num_inference_steps=req.num_inference_steps or 20
            )
        
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        buf.seek(0)
        return StreamingResponse(buf, media_type="image/png")
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)

# Keep your original echo for testing
class Prompt(BaseModel):
    text: str

@app.post("/echo")
def echo(p: Prompt):
    return {"length": len(p.text), "upper": p.text.upper()}

	import io, os
	from typing import Optional
	from fastapi import FastAPI
	from fastapi.responses import StreamingResponse, JSONResponse
	from pydantic import BaseModel
	from huggingface_hub import InferenceClient

	app = FastAPI(title="ImageGen API")

	# Choose a model that supports text->image via Inference API.
	# Using a model that works well with InferenceClient
	MODEL_ID = os.getenv("MODEL", "stabilityai/stable-diffusion-xl-base-1.0")
	HF_TOKEN = os.getenv("HF_TOKEN") # set in Space Secrets

	# Initialize client with or without token
	if HF_TOKEN:
	client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
	else:
	# Use public inference API (may have rate limits)
	client = InferenceClient(model=MODEL_ID)

	class GenReq(BaseModel):
	prompt: str
	width: int = 768
	height: int = 768
	seed: Optional[int] = None
	num_inference_steps: Optional[int] = None # some models ignore this

	@app.get("/")
	def root():
	return {"ok": True, "message": "API is live", "model": MODEL_ID}

	@app.post("/generate")
	def generate(req: GenReq):
	try:
	# Different models have different parameter requirements
	if "sd-turbo" in MODEL_ID.lower():
	# SD Turbo works with minimal parameters
	img = client.text_to_image(
	prompt=req.prompt,
	num_inference_steps=1, # SD Turbo works best with 1 step
	guidance_scale=0.0, # SD Turbo doesn't use guidance scale
	width=req.width,
	height=req.height
	)
	else:
	# Standard SDXL parameters
	img = client.text_to_image(
	prompt=req.prompt,
	width=req.width,
	height=req.height,
	seed=req.seed,
	num_inference_steps=req.num_inference_steps or 20
	)

	buf = io.BytesIO()
	img.save(buf, format="PNG")
	buf.seek(0)
	return StreamingResponse(buf, media_type="image/png")
	except Exception as e:
	return JSONResponse({"error": str(e)}, status_code=500)

	# Keep your original echo for testing
	class Prompt(BaseModel):
	text: str

	@app.post("/echo")
	def echo(p: Prompt):
	return {"length": len(p.text), "upper": p.text.upper()}