Upload 3 files

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.11-slim
+
+WORKDIR /code
+
+ENV HF_HOME /code/.cache
+
+RUN mkdir -p /code/.cache && chmod -R 777 /code/.cache
+
+COPY ./requirements.txt /code/requirements.txt
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY ./app.py /code/app.py
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,83 @@
+# 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 = 25
+
+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")
+    if not hf_token:
+        raise RuntimeError("HF_TOKEN environment variable not set!")
+
+    model_id = "rxmha125/sdxl-base-1.0-private" # <-- YOUR PRIVATE MODEL ID
+    print(f"Loading model: {model_id}")
+    
+    pipe = AutoPipelineForText2Image.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        variant="fp16",
+        use_safetensors=True,
+        token=hf_token
+    ).to("cuda")
+
+    # Optimization for speed and memory
+    pipe.enable_model_cpu_offload()
+    
+    app_state["pipe"] = pipe
+    print("Model loaded successfully.")
+    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 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:
+        # Generate the image
+        image = pipe(
+            prompt=request.prompt,
+            negative_prompt=request.negative_prompt,
+            num_inference_steps=request.steps
+        ).images[0]
+
+        # Convert image to Base64
+        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))

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+python-multipart
+diffusers>=0.25.0
+transformers
+accelerate
+torch --extra-index-url https://download.pytorch.org/whl/cu121