Workin on it

routers/training.py

@@ -13,14 +13,23 @@ import uuid
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionInpaintPipeline
 
 # tokenizer = AutoTokenizer.from_pretrained("openlm-research/open_llama_7b")
 # model = AutoModelForCausalLM.from_pretrained(
 #     "openlm-research/open_llama_7b", device_map="auto", load_in_4bit=True
 # )
 
-model_id_or_path = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
+# model_id_or_path = "runwayml/stable-diffusion-v1-5"
+# pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
+
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    revision="fp16",
+    torch_dtype=torch.float16,
+)
+
 pipe = pipe.to("cuda")
 
 
@@ -28,6 +37,7 @@ router = APIRouter()
 
 class ActionBody(BaseModel):
     url: str
+    maskUrl: str
     prompt: str
     strength: float
     guidance_scale: float
@@ -44,6 +54,12 @@ async def performAction(actionBody: ActionBody):
     response = requests.get(actionBody.url)
     init_image = Image.open(BytesIO(response.content)).convert("RGB")
     init_image = init_image.resize((actionBody.resizeW, actionBody.resizeH))
+
+    response = requests.get(actionBody.maskUrl)
+    init_image = Image.open(BytesIO(response.content)).convert("RGB")
+    init_image = init_image.resize((actionBody.resizeW, actionBody.resizeH))
+
+    
     images = pipe(prompt=actionBody.prompt, image=init_image, strength=actionBody.strength, guidance_scale=actionBody.guidance_scale).images
     print(images)
     buffered = BytesIO()