Update app.py

app.py

@@ -1,83 +1,85 @@
-# 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}")
+# 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}")
+    
+    # --- *** THIS IS THE CORRECTED PART *** ---
+    # We removed variant="fp16" and use_safetensors=True
+    # to load the available .bin files instead of the missing .safetensors.
+    pipe = AutoPipelineForText2Image.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16, # Keep for memory optimization
+        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))