Git Push Some update

Dockerfile

@@ -10,6 +10,9 @@ RUN apt-get update && apt-get install -y \
     python3-pip \
     && rm -rf /var/lib/apt/lists/*
 
+# Install accelerate for better GPU performance
+RUN pip3 install --no-cache-dir accelerate
+
 # Copy requirements first for better caching
 COPY --chown=user:user requirements.txt .
 RUN pip3 install --no-cache-dir -r requirements.txt
@@ -20,6 +23,10 @@ COPY --chown=user:user . .
 # Switch to non-root user
 USER user
 
+# Set environment variables for GPU
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+
 # Expose port
 EXPOSE 7860
 

app.py

@@ -1,129 +1,137 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-import torch
-from PIL import Image
-import io
-import base64
-from diffusers import StableDiffusionInpaintPipeline
-import gc
-from fastapi.responses import JSONResponse
-import logging
-
-app = FastAPI()
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Global variable for the model
-pipe = None
-
-def load_model():
-    global pipe
-    if pipe is None:
-        model_id = "Uminosachi/realisticVisionV51_v51VAE-inpainting"
-        try:
-            # Try CUDA first
-            if torch.cuda.is_available():
-                device = "cuda"
-                dtype = torch.float16
-            else:
-                # Fallback to CPU
-                device = "cpu"
-                dtype = torch.float32
-            
-            pipe = StableDiffusionInpaintPipeline.from_pretrained(
-                model_id,
-                torch_dtype=dtype,
-                safety_checker=None
-            ).to(device)
-            
-            if device == "cuda":
-                pipe.enable_attention_slicing(slice_size="max")
-                pipe.enable_sequential_cpu_offload()
-            
-            print(f"Model loaded on {device}")
-            
-        except Exception as e:
-            print(f"Error loading model: {str(e)}")
-            raise
-    return pipe
-
-@app.on_event("startup")
-async def startup_event():
-    try:
-        load_model()
-    except Exception as e:
-        print(f"Startup error: {str(e)}")
-
-def image_to_base64(image: Image.Image) -> str:
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode()
-
-@app.post("/inpaint")
-async def inpaint(
-    image: UploadFile = File(...),
-    mask: UploadFile = File(...),
-    prompt: str = "add some flowers and a fountain",
-    negative_prompt: str = "blurry, low quality, distorted"
-):
-    try:
-        # Add file size check (10MB limit)
-        max_size = 10 * 1024 * 1024  # 10MB
-        if len(await image.read()) > max_size or len(await mask.read()) > max_size:
-            return JSONResponse(
-                status_code=400,
-                content={"error": "File size too large. Maximum size is 10MB"}
-            )
-        
-        # Reset file positions
-        await image.seek(0)
-        await mask.seek(0)
-        
-        # Read and process input image
-        image_data = await image.read()
-        mask_data = await mask.read()
-        
-        original_image = Image.open(io.BytesIO(image_data))
-        mask_image = Image.open(io.BytesIO(mask_data))
-
-        # Resize to multiple of 8
-        width, height = (dim - dim % 8 for dim in original_image.size)
-        original_image = original_image.resize((width, height))
-        mask_image = mask_image.resize((width, height))
-        mask_image = mask_image.convert("L")
-
-        # Reduce steps for CPU
-        num_steps = 10 if not torch.cuda.is_available() else 20
-        
-        with torch.cuda.amp.autocast():
-            output_image = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=original_image,
-                mask_image=mask_image,
-                num_inference_steps=num_steps,  # Reduced steps
-                guidance_scale=7.5,
-            ).images[0]
-
-        # Convert output image to base64
-        output_base64 = image_to_base64(output_image)
-
-        # Clean up
-        torch.cuda.empty_cache()
-        gc.collect()
-
-        return {"status": "success", "image": output_base64}
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/health")
-async def health_check():
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+import torch
+from PIL import Image
+import io
+import base64
+from diffusers import StableDiffusionInpaintPipeline
+import gc
+from fastapi.responses import JSONResponse
+import logging
+
+app = FastAPI()
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global variable for the model
+pipe = None
+
+# Add max size limit
+MAX_SIZE = 512
+
+def load_model():
+    global pipe
+    if pipe is None:
+        model_id = "Uminosachi/realisticVisionV51_v51VAE-inpainting"
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.float16,
+                safety_checker=None
+            ).to(device)
+            
+            if device == "cuda":
+                pipe.enable_attention_slicing()
+            
+            print(f"Model loaded on {device} with optimizations")
+            
+        except Exception as e:
+            print(f"Error loading model: {str(e)}")
+            raise
+    return pipe
+
+@app.on_event("startup")
+async def startup_event():
+    try:
+        load_model()
+    except Exception as e:
+        print(f"Startup error: {str(e)}")
+
+def image_to_base64(image: Image.Image) -> str:
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
+def resize_for_condition_image(input_image: Image.Image, resolution: int):
+    input_width, input_height = input_image.size
+    aspect_ratio = input_height / input_width
+    
+    if input_height > input_width:
+        # vertical image
+        width = resolution
+        height = int(resolution * aspect_ratio)
+    else:
+        # horizontal image
+        height = resolution
+        width = int(resolution / aspect_ratio)
+    
+    return input_image.resize((width, height))
+
+@app.post("/inpaint")
+async def inpaint(
+    image: UploadFile = File(...),
+    mask: UploadFile = File(...),
+    prompt: str = "add some flowers and a fountain",
+    negative_prompt: str = "blurry, low quality, distorted"
+):
+    try:
+        # Add file size check (10MB limit)
+        max_size = 10 * 1024 * 1024  # 10MB
+        if len(await image.read()) > max_size or len(await mask.read()) > max_size:
+            return JSONResponse(
+                status_code=400,
+                content={"error": "File size too large. Maximum size is 10MB"}
+            )
+        
+        # Reset file positions
+        await image.seek(0)
+        await mask.seek(0)
+        
+        # Read and process input image
+        image_data = await image.read()
+        mask_data = await mask.read()
+        
+        original_image = Image.open(io.BytesIO(image_data))
+        mask_image = Image.open(io.BytesIO(mask_data))
+
+        # Resize images to smaller size
+        original_image = resize_for_condition_image(original_image, MAX_SIZE)
+        mask_image = resize_for_condition_image(mask_image, MAX_SIZE)
+        mask_image = mask_image.convert("L")
+
+        # Reduce steps even more for CPU
+        num_steps = 5 if not torch.cuda.is_available() else 20
+        
+        with torch.cuda.amp.autocast():
+            output_image = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image=original_image,
+                mask_image=mask_image,
+                num_inference_steps=num_steps,
+                guidance_scale=7.0,  # Slightly reduced for speed
+            ).images[0]
+
+        # Convert output image to base64
+        output_base64 = image_to_base64(output_image)
+
+        # Clean up
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        return {"status": "success", "image": output_base64}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
     return {"status": "healthy", "cuda_available": torch.cuda.is_available()} 

example_client.py

@@ -39,7 +39,7 @@ def main():
     # Replace these with your actual image paths
     image_path = r"C:\Users\M. Y\Downloads\t2.png"  # Replace with your image path
     mask_path = "generated_mask_1.png"   # Replace with your mask path
-    prompt = "add some flowers and a fountain"
+    prompt = "add some Chair tables and a fountain"
     
     result = call_inpaint_api(image_path, mask_path, prompt)
     if result: