Update app.py

app.py

@@ -12,7 +12,7 @@ import os
 class ImageRequest(BaseModel):
     prompt: str
     negative_prompt: str = ""
-    steps: int = 25
+    steps: int = 20 # Lowered for faster CPU inference
 
 class ImageResponse(BaseModel):
     image_base64: str
@@ -23,30 +23,26 @@ app_state = {}
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Load the model on startup
-    hf_token = os.getenv("HF_TOKEN")
+    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.")
 
-    # --- *** THIS IS THE ONLY LINE THAT CHANGES *** ---
-    # We now point directly to the original, public model repository.
-    model_id = "stabilityai/stable-diffusion-xl-base-1.0"
-    # --- ****************************************** ---
+    # --- *** 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}")
+    print(f"Loading model: {model_id} for CPU...")
     
+    # 2. Load the pipeline without GPU-specific settings
     pipe = AutoPipelineForText2Image.from_pretrained(
         model_id,
-        torch_dtype=torch.float16,
-        variant="fp16", # Use the optimized fp16 variant
-        use_safetensors=True,
         token=hf_token
-    ).to("cuda")
-
-    # Optimization for speed and memory
-    pipe.enable_model_cpu_offload()
+    )
+    # Note: We do not use .to("cuda")
+    # --- ************************************** ---
     
     app_state["pipe"] = pipe
-    print("Model loaded successfully.")
+    print("Model loaded successfully onto CPU.")
     yield
     # Clean up on shutdown
     app_state.clear()
@@ -57,7 +53,7 @@ app = FastAPI(lifespan=lifespan)
 
 @app.get("/")
 def root():
-    return {"status": "Text-to-Image API is running"}
+    return {"status": "Text-to-Image CPU API is running"}
 
 @app.post("/generate-image", response_model=ImageResponse)
 def generate_image(request: ImageRequest):