merge code

Dockerfile

@@ -35,4 +35,4 @@ USER appuser
 EXPOSE 7860
 
 # Run the server
-CMD ["python", "/app/intruct.py"]
+CMD ["python", "/app/merged_code.py"]

intruct.py

@@ -15,10 +15,32 @@ from io import BytesIO
 import os
 import base64
 from typing import List
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import StreamingResponse
+from PIL import Image, ImageDraw, ImageFilter
+import io
+import torch
+import numpy as np
+from diffusers import StableDiffusionInpaintPipeline
+import cv2
+
+
 
 # Initialize FastAPI app
 app = FastAPI()
 
+
+model_id_runway = "runwayml/stable-diffusion-inpainting"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+try:
+    pipe_runway = StableDiffusionInpaintPipeline.from_pretrained(model_id_runway)
+    pipe_runway.to(device)
+except Exception as e:
+    raise RuntimeError(f"Failed to load model: {e}")
+
+
+
 # Load the pre-trained InstructPix2Pix model for editing
 model_id = "timbrooks/instruct-pix2pix"
 pipe_edit = StableDiffusionInstructPix2PixPipeline.from_pretrained(
@@ -101,36 +123,6 @@ async def generate_image(prompt: str):
     except Exception as e:
         return {"error": str(e)}
 
-@app.get("/generate_multiple")
-async def generate_multiple_images(prompts: List[str]):
-    try:
-        # List to store base64-encoded images
-        generated_images = []
-        
-        # Generate an image for each prompt
-        for prompt in prompts:
-            image = pipe_generate(
-                prompt,
-                num_inference_steps=4,
-                guidance_scale=0
-            ).images[0]
-            
-            # Save image to buffer
-            buffer = BytesIO()
-            image.save(buffer, format="PNG")
-            buffer.seek(0)
-            
-            # Encode the image as base64
-            image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
-            generated_images.append({
-                "prompt": prompt,
-                "image_base64": image_base64
-            })
-        
-        return {"images": generated_images}
-    
-    except Exception as e:
-        return {"error": str(e)}
 
 @app.get("/health")
 async def health_check():
@@ -274,6 +266,196 @@ async def root():
     """
     return {"message": "InstructPix2Pix API is running. Use POST /edit-image/ or /inpaint/ to edit images."}
 
+
+
+# Helper functions
+def prepare_guided_image(original_image: Image, reference_image: Image, mask_image: Image) -> Image:
+    original_array = np.array(original_image)
+    reference_array = np.array(reference_image)
+    mask_array = np.array(mask_image) / 255.0
+    mask_array = mask_array[:, :, np.newaxis]
+    blended_array = original_array * (1 - mask_array) + reference_array * mask_array
+    return Image.fromarray(blended_array.astype(np.uint8))
+
+def soften_mask(mask_image: Image, softness: int = 5) -> Image:
+    from PIL import ImageFilter
+    return mask_image.filter(ImageFilter.GaussianBlur(radius=softness))
+
+def generate_rectangular_mask(image_size: tuple, x1: int = 100, y1: int = 100, x2: int = 200, y2: int = 200) -> Image:
+    mask = Image.new("L", image_size, 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([x1, y1, x2, y2], fill=255)
+    return mask
+
+def segment_tank(tank_image: Image) -> tuple[Image, Image]:
+    tank_array = np.array(tank_image.convert("RGB"))
+    tank_array = cv2.cvtColor(tank_array, cv2.COLOR_RGB2BGR)
+    hsv = cv2.cvtColor(tank_array, cv2.COLOR_BGR2HSV)
+    lower_snow = np.array([0, 0, 180])
+    upper_snow = np.array([180, 50, 255])
+    snow_mask = cv2.inRange(hsv, lower_snow, upper_snow)
+    tank_mask = cv2.bitwise_not(snow_mask)
+    kernel = np.ones((5, 5), np.uint8)
+    tank_mask = cv2.erode(tank_mask, kernel, iterations=1)
+    tank_mask = cv2.dilate(tank_mask, kernel, iterations=1)
+    tank_mask_image = Image.fromarray(tank_mask, mode="L")
+    tank_array_rgb = np.array(tank_image.convert("RGB"))
+    mask_array = tank_mask / 255.0
+    mask_array = mask_array[:, :, np.newaxis]
+    segmented_tank = (tank_array_rgb * mask_array).astype(np.uint8)
+    alpha = tank_mask
+    segmented_tank_rgba = np.zeros((tank_image.height, tank_image.width, 4), dtype=np.uint8)
+    segmented_tank_rgba[:, :, :3] = segmented_tank
+    segmented_tank_rgba[:, :, 3] = alpha
+    segmented_tank_image = Image.fromarray(segmented_tank_rgba, mode="RGBA")
+    return segmented_tank_image, tank_mask_image
+
+async def apply_camouflage_to_tank(tank_image: Image) -> Image:
+    segmented_tank, tank_mask = segment_tank(tank_image)
+    segmented_tank.save("segmented_tank.png")
+    tank_mask.save("tank_mask.png")
+    camouflaged_tank = pipe_runway(
+        prompt="Apply a grassy camouflage pattern with shades of green and brown to the tank, preserving its structure.",
+        image=segmented_tank.convert("RGB"),
+        mask_image=tank_mask,
+        strength=0.5,
+        guidance_scale=8.0,
+        num_inference_steps=50,
+        negative_prompt="snow, ice, rock, stone, boat, unrelated objects"
+    ).images[0]
+    camouflaged_tank_rgba = np.zeros((camouflaged_tank.height, camouflaged_tank.width, 4), dtype=np.uint8)
+    camouflaged_tank_rgba[:, :, :3] = np.array(camouflaged_tank)
+    camouflaged_tank_rgba[:, :, 3] = np.array(tank_mask)
+    camouflaged_tank_image = Image.fromarray(camouflaged_tank_rgba, mode="RGBA")
+    camouflaged_tank_image.save("camouflaged_tank.png")
+    return camouflaged_tank_image
+
+def fit_image_to_mask(original_image: Image, reference_image: Image, mask_x1: int, mask_y1: int, mask_x2: int, mask_y2: int) -> tuple:
+    mask_width = mask_x2 - mask_x1
+    mask_height = mask_y2 - mask_y1
+    if mask_width <= 0 or mask_height <= 0:
+        raise ValueError("Mask dimensions must be positive")
+    ref_width, ref_height = reference_image.size
+    aspect_ratio = ref_width / ref_height
+    if mask_width / mask_height > aspect_ratio:
+        new_height = mask_height
+        new_width = int(new_height * aspect_ratio)
+    else:
+        new_width = mask_width
+        new_height = int(new_width / aspect_ratio)
+    reference_image_resized = reference_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+    guided_image = original_image.copy().convert("RGB")
+    paste_x = mask_x1 + (mask_width - new_width) // 2
+    paste_y = mask_y1 + (mask_height - new_height) // 2
+    guided_image.paste(reference_image_resized, (paste_x, paste_y), reference_image_resized)
+    mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
+    return guided_image, mask_image
+
+# Endpoints
+@app.post("/inpaint/")
+async def inpaint_image(
+    image: UploadFile = File(...),
+    mask: UploadFile = File(...),
+    prompt: str = "Fill the masked area with appropriate content."
+):
+    try:
+        image_bytes = await image.read()
+        mask_bytes = await mask.read()
+        original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        mask_image = Image.open(io.BytesIO(mask_bytes)).convert("L")
+        if original_image.size != mask_image.size:
+            raise HTTPException(status_code=400, detail="Image and mask dimensions must match.")
+        result = pipe_runway(prompt=prompt, image=original_image, mask_image=mask_image).images[0]
+        result_bytes = io.BytesIO()
+        result.save(result_bytes, format="PNG")
+        result_bytes.seek(0)
+        return StreamingResponse(
+            result_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "attachment; filename=inpainted_image.png"}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during inpainting: {e}")
+
+@app.post("/inpaint-with-reference/")
+async def inpaint_with_reference(
+    image: UploadFile = File(...),
+    reference_image: UploadFile = File(...),
+    prompt: str = "Integrate the reference content naturally into the masked area, matching style and lighting.",
+    mask_x1: int = 100,
+    mask_y1: int = 100,
+    mask_x2: int = 200,
+    mask_y2: int = 200
+):
+    try:
+        image_bytes = await image.read()
+        reference_bytes = await reference_image.read()
+        original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        reference_image = Image.open(io.BytesIO(reference_bytes)).convert("RGB")
+        if original_image.size != reference_image.size:
+            reference_image = reference_image.resize(original_image.size, Image.Resampling.LANCZOS)
+        mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
+        softened_mask = soften_mask(mask_image, softness=5)
+        guided_image = prepare_guided_image(original_image, reference_image, softened_mask)
+        result = pipe_runway(
+            prompt=prompt,
+            image=guided_image,
+            mask_image=softened_mask,
+            strength=0.75,
+            guidance_scale=7.5
+        ).images[0]
+        result_bytes = io.BytesIO()
+        result.save(result_bytes, format="PNG")
+        result_bytes.seek(0)
+        return StreamingResponse(
+            result_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "attachment; filename=natural_inpaint_image.png"}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during natural inpainting: {e}")
+
+@app.post("/fit-image-to-mask/")
+async def fit_image_to_mask_endpoint(
+    image: UploadFile = File(...),
+    reference_image: UploadFile = File(...),
+    mask_x1: int = 200,
+    mask_y1: int = 200,
+    mask_x2: int = 500,
+    mask_y2: int = 500
+):
+    try:
+        image_bytes = await image.read()
+        reference_bytes = await reference_image.read()
+        original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        reference_image = Image.open(io.BytesIO(reference_bytes)).convert("RGB")
+        camouflaged_tank = await apply_camouflage_to_tank(reference_image)
+        guided_image, mask_image = fit_image_to_mask(original_image, camouflaged_tank, mask_x1, mask_y1, mask_x2, mask_y2)
+        guided_image.save("guided_image_before_blending.png")
+        softened_mask = soften_mask(mask_image, softness=2)
+        result = pipe_runway(
+            prompt="Blend the camouflaged tank into the grassy field with trees, ensuring a non-snowy environment, matching the style, lighting, and surroundings.",
+            image=guided_image,
+            mask_image=softened_mask,
+            strength=0.2,
+            guidance_scale=7.5,
+            num_inference_steps=50,
+            negative_prompt="snow, ice, rock, stone, boat, unrelated objects"
+        ).images[0]
+        result_bytes = io.BytesIO()
+        result.save(result_bytes, format="PNG")
+        result_bytes.seek(0)
+        return StreamingResponse(
+            result_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "attachment; filename=fitted_image.png"}
+        )
+    except ValueError as ve:
+        raise HTTPException(status_code=400, detail=f"ValueError in processing: {str(ve)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during fitting and inpainting: {str(e)}")
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)

merged_code.py

@@ -0,0 +1,461 @@
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import StreamingResponse
+import io
+import math
+from PIL import Image, ImageOps, ImageDraw
+import torch
+from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionInpaintPipeline
+from fastapi import FastAPI, Response
+from fastapi.responses import FileResponse
+import torch
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download, login
+from safetensors.torch import load_file
+from io import BytesIO
+import os
+import base64
+from typing import List
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import StreamingResponse
+from PIL import Image, ImageDraw, ImageFilter
+import io
+import torch
+import numpy as np
+from diffusers import StableDiffusionInpaintPipeline
+import cv2
+
+
+
+# Initialize FastAPI app
+app = FastAPI()
+
+
+model_id_runway = "runwayml/stable-diffusion-inpainting"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+try:
+    pipe_runway = StableDiffusionInpaintPipeline.from_pretrained(model_id_runway)
+    pipe_runway.to(device)
+except Exception as e:
+    raise RuntimeError(f"Failed to load model: {e}")
+
+
+
+# Load the pre-trained InstructPix2Pix model for editing
+model_id = "timbrooks/instruct-pix2pix"
+pipe_edit = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+    model_id, torch_dtype=torch.float16, safety_checker=None
+).to("cuda")
+
+# Load the pre-trained Inpainting model
+inpaint_model_id = "stabilityai/stable-diffusion-2-inpainting"
+pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+    inpaint_model_id, torch_dtype=torch.float16, safety_checker=None
+).to("cuda")
+
+# Default configuration values
+DEFAULT_STEPS = 50
+DEFAULT_TEXT_CFG = 7.5
+DEFAULT_IMAGE_CFG = 1.5
+DEFAULT_SEED = 1371
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+def load_model():
+    try:
+        # Login to Hugging Face if token is provided
+        if HF_TOKEN:
+            login(token=HF_TOKEN)
+            
+        base = "stabilityai/stable-diffusion-xl-base-1.0"
+        repo = "ByteDance/SDXL-Lightning"
+        ckpt = "sdxl_lightning_4step_unet.safetensors"
+
+        # Load model with explicit error handling
+        unet = UNet2DConditionModel.from_config(
+            base, 
+            subfolder="unet"
+        ).to("cuda", torch.float16)
+        
+        unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            base, 
+            unet=unet, 
+            torch_dtype=torch.float16, 
+            variant="fp16"
+        ).to("cuda")
+        
+        # Configure scheduler
+        pipe.scheduler = EulerDiscreteScheduler.from_config(
+            pipe.scheduler.config, 
+            timestep_spacing="trailing"
+        )
+        
+        return pipe
+    
+    except Exception as e:
+        raise Exception(f"Failed to load model: {str(e)}")
+
+# Load model at startup with error handling
+try:
+    pipe_generate = load_model()
+except Exception as e:
+    print(f"Model initialization failed: {str(e)}")
+    raise
+
+@app.get("/generate")
+async def generate_image(prompt: str):
+    try:
+        # Generate image
+        image = pipe_generate(
+            prompt,
+            num_inference_steps=4,
+            guidance_scale=0
+        ).images[0]
+        
+        # Save image to buffer
+        buffer = BytesIO()
+        image.save(buffer, format="PNG")
+        buffer.seek(0)
+        
+        return Response(content=buffer.getvalue(), media_type="image/png")
+    
+    except Exception as e:
+        return {"error": str(e)}
+
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+def process_image(input_image: Image.Image, instruction: str, steps: int, text_cfg_scale: float, image_cfg_scale: float, seed: int):
+    """
+    Process the input image with the given instruction using InstructPix2Pix.
+    """
+    # Resize image to fit model requirements
+    width, height = input_image.size
+    factor = 512 / max(width, height)
+    factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)
+    width = int((width * factor) // 64) * 64
+    height = int((height * factor) // 64) * 64
+    input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
+
+    if not instruction:
+        return input_image
+
+    # Set the random seed for reproducibility
+    generator = torch.manual_seed(seed)
+
+    # Generate the edited image
+    edited_image = pipe_edit(
+        instruction,
+        image=input_image,
+        guidance_scale=text_cfg_scale,
+        image_guidance_scale=image_cfg_scale,
+        num_inference_steps=steps,
+        generator=generator,
+    ).images[0]
+
+    return edited_image
+
+@app.post("/edit-image/")
+async def edit_image(
+    file: UploadFile = File(...),
+    instruction: str = Form(...),
+    steps: int = Form(default=DEFAULT_STEPS),
+    text_cfg_scale: float = Form(default=DEFAULT_TEXT_CFG),
+    image_cfg_scale: float = Form(default=DEFAULT_IMAGE_CFG),
+    seed: int = Form(default=DEFAULT_SEED)
+):
+    """
+    Endpoint to edit an image based on a text instruction.
+    """
+    # Read and convert the uploaded image
+    image_data = await file.read()
+    input_image = Image.open(io.BytesIO(image_data)).convert("RGB")
+
+    # Process the image
+    edited_image = process_image(input_image, instruction, steps, text_cfg_scale, image_cfg_scale, seed)
+
+    # Convert the edited image to bytes
+    img_byte_arr = io.BytesIO()
+    edited_image.save(img_byte_arr, format="PNG")
+    img_byte_arr.seek(0)
+
+    # Return the image as a streaming response
+    return StreamingResponse(img_byte_arr, media_type="image/png")
+
+# New endpoint for inpainting
+@app.post("/inpaint/")
+async def inpaint_image(
+    file: UploadFile = File(...),
+    prompt: str = Form(...),
+    mask_coordinates: str = Form(...),  # Format: "x1,y1,x2,y2" (top-left and bottom-right of the rectangle to inpaint)
+    steps: int = Form(default=DEFAULT_STEPS),
+    guidance_scale: float = Form(default=7.5),
+    seed: int = Form(default=DEFAULT_SEED)
+):
+    """
+    Endpoint to perform inpainting on an image.
+    - file: The input image to inpaint.
+    - prompt: The text prompt describing what to generate in the inpainted area.
+    - mask_coordinates: Coordinates of the rectangular area to inpaint (format: "x1,y1,x2,y2").
+    - steps: Number of inference steps.
+    - guidance_scale: Guidance scale for the inpainting process.
+    - seed: Random seed for reproducibility.
+    """
+    try:
+        # Read and convert the uploaded image
+        image_data = await file.read()
+        input_image = Image.open(io.BytesIO(image_data)).convert("RGB")
+
+        # Resize image to fit model requirements (must be divisible by 8 for inpainting)
+        width, height = input_image.size
+        factor = 512 / max(width, height)
+        factor = math.ceil(min(width, height) * factor / 8) * 8 / min(width, height)
+        width = int((width * factor) // 8) * 8
+        height = int((height * factor) // 8) * 8
+        input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
+
+        # Create a mask for inpainting
+        mask = Image.new("L", (width, height), 0)  # Black image (0 = no inpainting)
+        draw = ImageDraw.Draw(mask)
+
+        # Parse the mask coordinates
+        try:
+            x1, y1, x2, y2 = map(int, mask_coordinates.split(","))
+            # Adjust coordinates based on resized image
+            x1 = int(x1 * factor)
+            y1 = int(y1 * factor)
+            x2 = int(x2 * factor)
+            y2 = int(y2 * factor)
+        except ValueError:
+            return {"error": "Invalid mask coordinates format. Use 'x1,y1,x2,y2'."}
+
+        # Draw a white rectangle on the mask (255 = area to inpaint)
+        draw.rectangle([x1, y1, x2, y2], fill=255)
+
+        # Set the random seed for reproducibility
+        generator = torch.manual_seed(seed)
+
+        # Perform inpainting
+        inpainted_image = pipe_inpaint(
+            prompt=prompt,
+            image=input_image,
+            mask_image=mask,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        ).images[0]
+
+        # Convert the inpainted image to bytes
+        img_byte_arr = io.BytesIO()
+        inpainted_image.save(img_byte_arr, format="PNG")
+        img_byte_arr.seek(0)
+
+        # Return the image as a streaming response
+        return StreamingResponse(img_byte_arr, media_type="image/png")
+
+    except Exception as e:
+        return {"error": str(e)}
+
+@app.get("/")
+async def root():
+    """
+    Root endpoint for basic health check.
+    """
+    return {"message": "InstructPix2Pix API is running. Use POST /edit-image/ or /inpaint/ to edit images."}
+
+
+
+# Helper functions
+def prepare_guided_image(original_image: Image, reference_image: Image, mask_image: Image) -> Image:
+    original_array = np.array(original_image)
+    reference_array = np.array(reference_image)
+    mask_array = np.array(mask_image) / 255.0
+    mask_array = mask_array[:, :, np.newaxis]
+    blended_array = original_array * (1 - mask_array) + reference_array * mask_array
+    return Image.fromarray(blended_array.astype(np.uint8))
+
+def soften_mask(mask_image: Image, softness: int = 5) -> Image:
+    from PIL import ImageFilter
+    return mask_image.filter(ImageFilter.GaussianBlur(radius=softness))
+
+def generate_rectangular_mask(image_size: tuple, x1: int = 100, y1: int = 100, x2: int = 200, y2: int = 200) -> Image:
+    mask = Image.new("L", image_size, 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([x1, y1, x2, y2], fill=255)
+    return mask
+
+def segment_tank(tank_image: Image) -> tuple[Image, Image]:
+    tank_array = np.array(tank_image.convert("RGB"))
+    tank_array = cv2.cvtColor(tank_array, cv2.COLOR_RGB2BGR)
+    hsv = cv2.cvtColor(tank_array, cv2.COLOR_BGR2HSV)
+    lower_snow = np.array([0, 0, 180])
+    upper_snow = np.array([180, 50, 255])
+    snow_mask = cv2.inRange(hsv, lower_snow, upper_snow)
+    tank_mask = cv2.bitwise_not(snow_mask)
+    kernel = np.ones((5, 5), np.uint8)
+    tank_mask = cv2.erode(tank_mask, kernel, iterations=1)
+    tank_mask = cv2.dilate(tank_mask, kernel, iterations=1)
+    tank_mask_image = Image.fromarray(tank_mask, mode="L")
+    tank_array_rgb = np.array(tank_image.convert("RGB"))
+    mask_array = tank_mask / 255.0
+    mask_array = mask_array[:, :, np.newaxis]
+    segmented_tank = (tank_array_rgb * mask_array).astype(np.uint8)
+    alpha = tank_mask
+    segmented_tank_rgba = np.zeros((tank_image.height, tank_image.width, 4), dtype=np.uint8)
+    segmented_tank_rgba[:, :, :3] = segmented_tank
+    segmented_tank_rgba[:, :, 3] = alpha
+    segmented_tank_image = Image.fromarray(segmented_tank_rgba, mode="RGBA")
+    return segmented_tank_image, tank_mask_image
+
+async def apply_camouflage_to_tank(tank_image: Image) -> Image:
+    segmented_tank, tank_mask = segment_tank(tank_image)
+    segmented_tank.save("segmented_tank.png")
+    tank_mask.save("tank_mask.png")
+    camouflaged_tank = pipe_runway(
+        prompt="Apply a grassy camouflage pattern with shades of green and brown to the tank, preserving its structure.",
+        image=segmented_tank.convert("RGB"),
+        mask_image=tank_mask,
+        strength=0.5,
+        guidance_scale=8.0,
+        num_inference_steps=50,
+        negative_prompt="snow, ice, rock, stone, boat, unrelated objects"
+    ).images[0]
+    camouflaged_tank_rgba = np.zeros((camouflaged_tank.height, camouflaged_tank.width, 4), dtype=np.uint8)
+    camouflaged_tank_rgba[:, :, :3] = np.array(camouflaged_tank)
+    camouflaged_tank_rgba[:, :, 3] = np.array(tank_mask)
+    camouflaged_tank_image = Image.fromarray(camouflaged_tank_rgba, mode="RGBA")
+    camouflaged_tank_image.save("camouflaged_tank.png")
+    return camouflaged_tank_image
+
+def fit_image_to_mask(original_image: Image, reference_image: Image, mask_x1: int, mask_y1: int, mask_x2: int, mask_y2: int) -> tuple:
+    mask_width = mask_x2 - mask_x1
+    mask_height = mask_y2 - mask_y1
+    if mask_width <= 0 or mask_height <= 0:
+        raise ValueError("Mask dimensions must be positive")
+    ref_width, ref_height = reference_image.size
+    aspect_ratio = ref_width / ref_height
+    if mask_width / mask_height > aspect_ratio:
+        new_height = mask_height
+        new_width = int(new_height * aspect_ratio)
+    else:
+        new_width = mask_width
+        new_height = int(new_width / aspect_ratio)
+    reference_image_resized = reference_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+    guided_image = original_image.copy().convert("RGB")
+    paste_x = mask_x1 + (mask_width - new_width) // 2
+    paste_y = mask_y1 + (mask_height - new_height) // 2
+    guided_image.paste(reference_image_resized, (paste_x, paste_y), reference_image_resized)
+    mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
+    return guided_image, mask_image
+
+# Endpoints
+@app.post("/inpaint/")
+async def inpaint_image(
+    image: UploadFile = File(...),
+    mask: UploadFile = File(...),
+    prompt: str = "Fill the masked area with appropriate content."
+):
+    try:
+        image_bytes = await image.read()
+        mask_bytes = await mask.read()
+        original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        mask_image = Image.open(io.BytesIO(mask_bytes)).convert("L")
+        if original_image.size != mask_image.size:
+            raise HTTPException(status_code=400, detail="Image and mask dimensions must match.")
+        result = pipe_runway(prompt=prompt, image=original_image, mask_image=mask_image).images[0]
+        result_bytes = io.BytesIO()
+        result.save(result_bytes, format="PNG")
+        result_bytes.seek(0)
+        return StreamingResponse(
+            result_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "attachment; filename=inpainted_image.png"}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during inpainting: {e}")
+
+@app.post("/inpaint-with-reference/")
+async def inpaint_with_reference(
+    image: UploadFile = File(...),
+    reference_image: UploadFile = File(...),
+    prompt: str = "Integrate the reference content naturally into the masked area, matching style and lighting.",
+    mask_x1: int = 100,
+    mask_y1: int = 100,
+    mask_x2: int = 200,
+    mask_y2: int = 200
+):
+    try:
+        image_bytes = await image.read()
+        reference_bytes = await reference_image.read()
+        original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        reference_image = Image.open(io.BytesIO(reference_bytes)).convert("RGB")
+        if original_image.size != reference_image.size:
+            reference_image = reference_image.resize(original_image.size, Image.Resampling.LANCZOS)
+        mask_image = generate_rectangular_mask(original_image.size, mask_x1, mask_y1, mask_x2, mask_y2)
+        softened_mask = soften_mask(mask_image, softness=5)
+        guided_image = prepare_guided_image(original_image, reference_image, softened_mask)
+        result = pipe_runway(
+            prompt=prompt,
+            image=guided_image,
+            mask_image=softened_mask,
+            strength=0.75,
+            guidance_scale=7.5
+        ).images[0]
+        result_bytes = io.BytesIO()
+        result.save(result_bytes, format="PNG")
+        result_bytes.seek(0)
+        return StreamingResponse(
+            result_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "attachment; filename=natural_inpaint_image.png"}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during natural inpainting: {e}")
+
+@app.post("/fit-image-to-mask/")
+async def fit_image_to_mask_endpoint(
+    image: UploadFile = File(...),
+    reference_image: UploadFile = File(...),
+    mask_x1: int = 200,
+    mask_y1: int = 200,
+    mask_x2: int = 500,
+    mask_y2: int = 500
+):
+    try:
+        image_bytes = await image.read()
+        reference_bytes = await reference_image.read()
+        original_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        reference_image = Image.open(io.BytesIO(reference_bytes)).convert("RGB")
+        camouflaged_tank = await apply_camouflage_to_tank(reference_image)
+        guided_image, mask_image = fit_image_to_mask(original_image, camouflaged_tank, mask_x1, mask_y1, mask_x2, mask_y2)
+        guided_image.save("guided_image_before_blending.png")
+        softened_mask = soften_mask(mask_image, softness=2)
+        result = pipe_runway(
+            prompt="Blend the camouflaged tank into the grassy field with trees, ensuring a non-snowy environment, matching the style, lighting, and surroundings.",
+            image=guided_image,
+            mask_image=softened_mask,
+            strength=0.2,
+            guidance_scale=7.5,
+            num_inference_steps=50,
+            negative_prompt="snow, ice, rock, stone, boat, unrelated objects"
+        ).images[0]
+        result_bytes = io.BytesIO()
+        result.save(result_bytes, format="PNG")
+        result_bytes.seek(0)
+        return StreamingResponse(
+            result_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": "attachment; filename=fitted_image.png"}
+        )
+    except ValueError as ve:
+        raise HTTPException(status_code=400, detail=f"ValueError in processing: {str(ve)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during fitting and inpainting: {str(e)}")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -8,4 +8,6 @@ transformers
 pillow
 numpy
 blenderproc
-python-multipart
+python-multipart
+opencv-python
+opencv-python-headless