Upload app.py

app.py

@@ -1,99 +1,40 @@
-# app.py
-import os
-import io
-import base64
-import logging
-from typing import Optional
-
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from PIL import Image
-import torch
-
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
 from diffusers import DiffusionPipeline
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("ssd_service")
-
-class GenerateRequest(BaseModel):
-    prompt: str
-    height: Optional[int] = None
-    width: Optional[int] = None
-    num_inference_steps: Optional[int] = 25
-
-app = FastAPI(title="Diffusers SSD-1B service")
-
-# Lazy global
-PIPELINE = None
-DEVICE = "cpu"
-
-def load_model():
-    global PIPELINE, DEVICE
-    # Prefer CUDA if available
-    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    logger.info(f"Loading model to device: {DEVICE}")
-
-    # If you need authentication for private models, set HUGGINGFACE_TOKEN env var
-    # hf_token = os.environ.get("HUGGINGFACE_TOKEN", None)
-    # auth = {"token": hf_token} if hf_token else None
-
-    # This can take time — load once on startup
-    PIPELINE = DiffusionPipeline.from_pretrained(
-        "segmind/SSD-1B",
-        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-        use_safetensors=True,
-        # If you want to use accelerate or offload, consider integration here
-    )
-
-    # move to device
-    PIPELINE = PIPELINE.to(DEVICE)
-
-    # Optionally enable attention slicing to save memory
-    try:
-        PIPELINE.enable_attention_slicing()
-    except Exception:
-        pass
-
-@app.on_event("startup")
-def startup_event():
-    load_model()
-    logger.info("Model loaded and service startup finished.")
-
-@app.get("/health")
-def health():
-    return {"status": "ok", "device": DEVICE}
-
-@app.post("/generate")
-def generate(req: GenerateRequest):
-    global PIPELINE
-    if PIPELINE is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-
-    prompt = req.prompt
-    if not prompt:
-        raise HTTPException(status_code=400, detail="prompt is required")
-
-    # Run inference (synchronous). Consider queuing for heavy load.
-    try:
-        logger.info(f"Generating image for prompt: {prompt[:80]}")
-        # Customize call as required; return PIL image
-        output = PIPELINE(prompt=prompt, num_inference_steps=req.num_inference_steps)
-        images = output.images
-        if not images:
-            raise HTTPException(status_code=500, detail="No image produced")
-
-        img = images[0]
-
-        # Resize if requested
-        if req.width and req.height:
-            img = img.resize((req.width, req.height), Image.LANCZOS)
-
-        # Convert to PNG bytes
-        buf = io.BytesIO()
-        img.save(buf, format="PNG")
-        buf.seek(0)
-        b64 = base64.b64encode(buf.read()).decode("utf-8")
-        return {"image_base64": b64}
-    except Exception as e:
-        logger.exception("Generation failed")
-        raise HTTPException(status_code=500, detail=str(e))
+import torch
+import base64
+from io import BytesIO
+
+# create fastapi instance
+app = FastAPI()
+
+# Detect device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load the correct pipeline
+pipe = DiffusionPipeline.from_pretrained(
+    "segmind/SSD-1B",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    use_safetensors=True,
+)
+pipe = pipe.to(device)
+
+# Optional: reduce VRAM usage
+pipe.enable_attention_slicing()
+
+@app.get("/")
+def home():
+    return {"message": "Segmind SSD-1B API running"}
+
+# Define function to generate image from text prompt
+@app.post("/generate-image/")
+def generate_image(prompt: str, negative_prompt: str = None):
+    # Run inference
+    image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]
+    
+    # Convert to base64 string
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    return JSONResponse(content={"image_base64": img_str})