creek
The AI Image Editor of 2030git clone https://github.com/maydayv7/creek
models/inpainting.py (8353B)
1"""
2Inpainting
3Original file: https://colab.research.google.com/drive/15HZcoDGdzTaDO9d1uIUMI8LQToCP4ryC
4"""
5
6## Fal.ai API
7
8!pip install fal-client
9
10import numpy as np
11import pandas as pd
12import fal_client
13import requests
14import os
15import numpy as np
16import scipy.ndimage
17from io import BytesIO
18from PIL import Image, ImageFilter
19from IPython.display import display
20
21os.environ["FAL_KEY"] = "YOUR_API_KEY"
22
23CLEAN_IMAGE_PATH = "clean_image_path"
24DRAWN_IMAGE_PATH = "drawn_image_path"
25MASK_OUTPUT_PATH = "solid_mask.png"
26
27
28def generate_and_save_mask(clean_path, drawn_path, save_path):
29 print("Generating mask with robust filtering...")
30
31 if not os.path.exists(clean_path) or not os.path.exists(drawn_path):
32 print(" Error: Input images not found.")
33 return False
34
35 img_clean = Image.open(clean_path).convert("RGB").resize((1024, 1024))
36 img_drawn = Image.open(drawn_path).convert("RGB").resize((1024, 1024))
37
38 clean_blur = np.array(img_clean.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16)
39 drawn_blur = np.array(img_drawn.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16)
40
41 diff_arr = np.abs(drawn_blur - clean_blur)
42 mask_arr = np.max(diff_arr, axis=2)
43
44 mask_binary = mask_arr > 30
45
46 mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
47 mask = Image.fromarray((mask_filled * 255).astype(np.uint8))
48 mask = mask.filter(ImageFilter.MaxFilter(9))
49
50 mask.save(save_path)
51 img_clean.save("resized_clean.png")
52
53 print(" The AI will ONLY edit the WHITE area below:")
54 display(mask.resize((256, 256)))
55
56 return True
57
58
59def run_fal_inpainting(image_path, mask_path):
60 print("Uploading images to Fal.ai...")
61
62 image_url = fal_client.upload_file(image_path)
63 mask_url = fal_client.upload_file(mask_path)
64
65 print("Running Flux Dev Fill...")
66
67 try:
68 handler = fal_client.submit(
69 "fal-ai/flux-lora-fill",
70 arguments={
71 "prompt": "The image shows a river running through a lush green valley surrounded by trees, plants, grass, and poles. In the background, the sky is filled with clouds, creating a peaceful atmosphere.",
72 "image_url": image_url,
73 "mask_url": mask_url,
74 "guidance_scale": 30,
75 "num_inference_steps": 28,
76 "enable_safety_checker": False,
77 },
78 )
79
80 result = handler.get()
81
82 if "images" in result:
83 output_url = result["images"][0]["url"]
84 print(f"Success! Image generated: {output_url}")
85
86 response = requests.get(output_url)
87 img = Image.open(BytesIO(response.content))
88 img.save("final_output.png")
89 display(img)
90 else:
91 print(" API returned no images.")
92 print(result)
93
94 except Exception as e:
95 print(f"Error during API call: {e}")
96
97
98# MAIN EXECUTION
99
100if __name__ == "__main__":
101 success = generate_and_save_mask(
102 CLEAN_IMAGE_PATH, DRAWN_IMAGE_PATH, MASK_OUTPUT_PATH
103 )
104 if success:
105 run_fal_inpainting("resized_clean.png", MASK_OUTPUT_PATH)
106
107## Stable Diffusion
108
109!pip install torch diffusers transformers accelerate gradio
110
111import torch
112import numpy as np
113import matplotlib.pyplot as plt
114from PIL import Image, ImageFilter
115from diffusers import AutoPipelineForInpainting
116import os
117import scipy.ndimage
118
119
120def load_model():
121 model_id = "runwayml/stable-diffusion-inpainting"
122 device = "cuda" if torch.cuda.is_available() else "cpu"
123
124 print(f"Loading model to {device} (this may take a minute)...")
125 try:
126 pipe = AutoPipelineForInpainting.from_pretrained(
127 model_id, torch_dtype=torch.float16, variant="fp16"
128 ).to(device)
129 pipe.enable_attention_slicing()
130 return pipe
131 except Exception as e:
132 print(f"Error loading model: {e}")
133 return None
134
135
136def run_sketch_to_image(
137 pipe, clean_path, drawn_path, prompt, negative_prompt="", strength=0.85, seed=42
138):
139 if not os.path.exists(clean_path) or not os.path.exists(drawn_path):
140 print(" Error: Images not found.")
141 return
142
143 img_clean = Image.open(clean_path).convert("RGB").resize((512, 512))
144 img_drawn = Image.open(drawn_path).convert("RGB").resize((512, 512))
145
146 clean_blur = np.array(
147 img_clean.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16
148 )
149 drawn_blur = np.array(
150 img_drawn.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16
151 )
152
153 diff_arr = np.abs(drawn_blur - clean_blur)
154
155 mask_arr = np.max(diff_arr, axis=2)
156
157 mask_binary = mask_arr > 30
158
159 mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
160 mask = Image.fromarray((mask_filled * 255).astype(np.uint8))
161
162 mask = mask.filter(ImageFilter.MaxFilter(9))
163
164 generator = torch.Generator(device="cuda").manual_seed(seed)
165
166 print(f" Generating with Strength {strength}...")
167 with torch.inference_mode():
168 output = pipe(
169 prompt=prompt,
170 negative_prompt=negative_prompt,
171 image=img_drawn,
172 mask_image=mask,
173 strength=strength,
174 guidance_scale=8.0,
175 num_inference_steps=50,
176 generator=generator,
177 ).images[0]
178
179 fig, axs = plt.subplots(1, 4, figsize=(20, 6))
180 axs[0].imshow(img_clean)
181 axs[0].set_title("Initial")
182 axs[1].imshow(img_drawn)
183 axs[1].set_title("Sketch")
184 axs[2].imshow(mask, cmap="gray")
185 axs[2].set_title("Calculated Mask (Fixed)")
186 axs[3].imshow(output)
187 axs[3].set_title("Result")
188 plt.tight_layout()
189 plt.show()
190 return output
191
192
193if __name__ == "__main__":
194 if "pipeline" not in globals():
195 pipeline = load_model()
196
197 if pipeline:
198 run_sketch_to_image(
199 pipe=pipeline,
200 clean_path="/kaggle/input/sample/scenery.jpg.png",
201 drawn_path="/kaggle/input/sample/scenery_sketch.jpg",
202 prompt="The image shows a river running through a lush green valley surrounded by trees, plants, grass, and poles. In the background, the sky is filled with clouds, creating a peaceful atmosphere.",
203 negative_prompt="",
204 strength=0.85,
205 seed=100,
206 )
207
208
209## Captioning
210
211!pip install -q einops timm
212
213import torch
214from transformers import AutoProcessor, AutoModelForCausalLM
215from PIL import Image
216import time
217
218device = "cuda" if torch.cuda.is_available() else "cpu"
219model_id = "microsoft/Florence-2-base"
220
221print(f" Loading Model: {model_id}...")
222
223model = AutoModelForCausalLM.from_pretrained(
224 model_id, trust_remote_code=True, torch_dtype=torch.float16
225).to(device)
226
227processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
228
229print("Model Loaded & Moved to GPU. Ready for Inference.")
230
231
232def generate_concise_caption(image_path):
233 device = "cuda" if torch.cuda.is_available() else "cpu"
234 dtype = torch.float16
235
236 torch.cuda.reset_peak_memory_stats()
237 torch.cuda.synchronize()
238 start_time = time.time()
239
240 image = Image.open(image_path).convert("RGB")
241
242 prompt_task = "<DETAILED_CAPTION>"
243
244 inputs = processor(text=prompt_task, images=image, return_tensors="pt").to(
245 device, dtype
246 )
247
248 generated_ids = model.generate(
249 input_ids=inputs["input_ids"],
250 pixel_values=inputs["pixel_values"],
251 max_new_tokens=64,
252 num_beams=1,
253 do_sample=False,
254 )
255
256 generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
257 caption = processor.post_process_generation(
258 generated_text, task=prompt_task, image_size=(image.width, image.height)
259 )[prompt_task]
260
261 caption = caption.strip()
262 if not caption.endswith("."):
263 last_period_index = caption.rfind(".")
264 if last_period_index != -1:
265 caption = caption[: last_period_index + 1]
266
267 torch.cuda.synchronize()
268 end_time = time.time()
269 latency = end_time - start_time
270 mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
271
272 print("-" * 30)
273 print(f"Result: {caption}")
274 print(f"Tokens: ~{len(caption.split())} words") # Approx count
275 print("-" * 30)
276 print(f"⚡ Latency: {latency:.4f}s")
277 print(f"VRAM: {mem_gb:.2f} GB")
278 print("-" * 30)
279
280 return caption
281
282
283text_output = generate_concise_caption("/kaggle/input/test1/tree_sketch.jpg")