Ragendom/SnapChefAI

1import os2 3# Set safe writable directories for model cache4os.environ["HF_HOME"] = "/data/.cache"5os.environ["TRANSFORMERS_CACHE"] = "/data/.cache"6 7import gradio as gr8import torch9import numpy as np10from PIL import Image11import base6412from io import BytesIO13from diffusers import StableDiffusionInpaintPipeline14 15# Load the pipeline once globally16pipe = StableDiffusionInpaintPipeline.from_pretrained(17    "stabilityai/stable-diffusion-2-inpainting",18    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float3219)20pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")21 22def decode_base64_image(data_url):23    if data_url.startswith("data:image"):24        base64_data = data_url.split(",")[1]25    else:26        base64_data = data_url27    return Image.open(BytesIO(base64.b64decode(base64_data)))28 29def generate_image(init_image_b64, mask_image_b64, prompt):30    init_image = decode_base64_image(init_image_b64).convert("RGB").resize((512, 512))31    mask_image = decode_base64_image(mask_image_b64).convert("L").resize((512, 512))32    33    result = pipe(34        prompt=prompt,35        image=init_image,36        mask_image=mask_image37    ).images[0]38    39    buffered = BytesIO()40    result.save(buffered, format="PNG")41    return "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode()42 43# Gradio input + output config44iface = gr.Interface(45    fn=generate_image,46    inputs=[47        gr.Text(label="Init Image (base64 PNG)"),48        gr.Text(label="Mask Image (base64 PNG)"),49        gr.Text(label="Prompt")50    ],51    outputs=gr.Text(label="Generated Image (base64 PNG)"),52    title="SnapChefAI - SD Inpainting API"53)54 55# Launch it!56iface.launch()57