Upload 5 files

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,13 @@
----
-title: Space1
-emoji: 🐨
-colorFrom: purple
-colorTo: indigo
-sdk: streamlit
-sdk_version: 1.39.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Stable Edit
+emoji: 📊
+colorFrom: indigo
+colorTo: indigo
+sdk: streamlit
+sdk_version: 1.21.0
+app_file: app.py
+pinned: false
+license: cc
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,138 @@
+import numpy as np
+import pandas as pd
+from PIL import Image
+from collections import defaultdict
+
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+
+import matplotlib as mpl
+
+from model import device, segment_image, inpaint
+
+
+# define utils and helpers
+def closest_number(n, m=8):
+    """ Obtains closest number to n that is divisble by m """
+    return int(n/m) * m
+
+
+def get_mask_from_rectangles(image, mask, height, width, drawing_mode='rect'):
+    # Create a canvas component
+    canvas_result = st_canvas(
+        fill_color="rgba(255, 165, 0, 0.3)",  
+        stroke_width=2,
+        stroke_color="#000000", 
+        background_image=image,
+        update_streamlit=True,
+        height=height,
+        width=width,
+        drawing_mode=drawing_mode,
+        point_display_radius=5,
+        key="canvas",
+    )
+
+    # get selections from mask
+    if canvas_result.json_data is not None:
+        objects = pd.json_normalize(canvas_result.json_data["objects"])
+        for col in objects.select_dtypes(include=["object"]).columns:
+            objects[col] = objects[col].astype("str")
+
+        if len(objects) > 0:
+            left_coords = objects.left.to_numpy()
+            top_coords = objects.top.to_numpy()
+            right_coords = left_coords + objects.width.to_numpy()
+            bottom_coords = top_coords + objects.height.to_numpy()
+
+            # add selections to mask
+            for (left, top, right, bottom) in zip(left_coords, top_coords, right_coords, bottom_coords):
+                cropped = image.crop((left, top, right, bottom))
+                st.image(cropped)
+                mask[top:bottom, left:right] = 255
+
+            st.header("Mask Created!")
+            st.image(mask)
+
+    return mask
+
+
+def get_mask(image, edit_method, height, width):
+    mask = np.zeros((height, width), dtype=np.uint8)
+
+    if edit_method == "AutoSegment Area":
+
+        # get displayable segmented image
+        seg_prediction, segment_labels = segment_image(image)
+        seg = seg_prediction['segmentation'].cpu().numpy()
+        viridis = mpl.colormaps.get_cmap('viridis').resampled(np.max(seg))
+        seg_image = Image.fromarray(np.uint8(viridis(seg)*255))
+
+        st.image(seg_image)
+
+        # prompt user to select valid labels to edit
+        seg_selections = st.multiselect("Choose segments", zip(segment_labels.keys(), segment_labels.values()))
+        if seg_selections:
+            tgts = []
+            for s in seg_selections:
+                tgts.append(s[0])
+
+            mask = Image.fromarray(np.array([(seg == t) for t in tgts]).sum(axis=0).astype(np.uint8)*255)
+            st.header("Mask Created!")
+            st.image(mask)
+
+    elif edit_method == "Draw Custom Area":
+        mask = get_mask_from_rectangles(image, mask, height, width)
+
+
+    return mask
+
+
+
+if __name__ == '__main__':
+    
+    st.title("Stable Edit")
+    st.title("Edit your photos with Stable Diffusion!")
+
+    st.write(f"Device found: {device}")
+
+    sf = st.text_input("Please enter resizing scale factor to downsize image (default=2)", value="2")
+    try: 
+        sf = int(sf)
+    except:
+        sf.write("Error with input scale factor, setting to default value of 2, please re-enter above to change it")
+        sf = 2
+
+    # upload image
+    filename = st.file_uploader("upload an image")
+
+    if filename:
+        image = Image.open(filename)
+
+        width, height = image.size
+        width, height = closest_number(width/sf), closest_number(height/sf)
+        image = image.resize((width, height))
+
+        st.image(image)
+        # st.write(f"{width} {height}")
+
+        # Select an editing method
+        edit_method = st.selectbox("Select Edit Method", ("AutoSegment Area", "Draw Custom Area"))
+
+        if edit_method:
+            mask = get_mask(image, edit_method, height, width)
+
+            # get inpainted images
+            prompt = st.text_input("Please enter prompt for image inpainting", value="")
+
+            if prompt: #  and isinstance(seed, int):
+                st.write("Inpainting Images, patience is a virtue and this will take a while to run on a CPU :)")
+                images = inpaint(image, mask, width, height, prompt=prompt, seed=0, guidance_scale=17.5, num_samples=3)
+
+                # display all images
+                st.write("Original Image")
+                st.image(image)
+                for i, img in enumerate(images, 1):
+                    st.write(f"result: {i}")
+                    st.image(img)
+
+    

model.py

@@ -0,0 +1,86 @@
+import torch
+from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+from diffusers import StableDiffusionInpaintPipeline # , DiffusionPipeline
+
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+# Image segmentation
+seg_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+seg_model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+
+def segment_image(image):
+    inputs = seg_processor(image, return_tensors="pt")
+    
+    with torch.no_grad():
+        seg_outputs = seg_model(**inputs)
+
+    # get prediction dict
+    seg_prediction = seg_processor.post_process_panoptic_segmentation(seg_outputs, target_sizes=[image.size[::-1]])[0]
+
+    # get segment labels dict
+    segment_labels = {}
+    for segment in seg_prediction['segments_info']:
+        segment_id = segment['id']
+        segment_label_id = segment['label_id']
+        segment_label = seg_model.config.id2label[segment_label_id]
+
+        segment_labels.update({segment_id : segment_label})
+
+    return seg_prediction, segment_labels
+
+# Image inpainting pipeline
+# get Stable Diffusion model for image inpainting
+if device == 'cuda':
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting",
+        torch_dtype=torch.float16,
+    ).to(device)
+else:
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting",
+        torch_dtype=torch.bfloat16, 
+        device_map="auto"
+    )
+
+# pipe = StableDiffusionInpaintPipeline.from_pretrained( # DiffusionPipeline.from_pretrained(
+#         "runwayml/stable-diffusion-inpainting",
+#         revision="fp16",
+#         torch_dtype=torch.bfloat16, 
+#         # device_map="auto" # use for Hugging face spaces
+#     )
+# pipe.to(device) # use for local environment
+
+def inpaint(image, mask, W, H, prompt="", seed=0, guidance_scale=17.5, num_samples=3):
+    """ Uses Stable Diffusion model to inpaint image 
+        Inputs: 
+            image - input image (PIL or torch tensor)
+            mask - mask for inpainting same size as image (PIL or troch tensor)
+            W - size of image
+            H - size of mask
+            prompt - prompt for inpainting
+            seed - random seed 
+        Outputs:
+            images - output images
+        """
+    generator = torch.Generator(device=device).manual_seed(seed)
+    images = pipe(
+        prompt=prompt,
+        image=image,
+        mask_image=mask, # ensure mask is same type as image
+        height=H,
+        width=W,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        num_images_per_prompt=num_samples,
+    ).images
+    
+    return images
+
+    
+
+    
+
+
+    

requirements.txt

@@ -0,0 +1,9 @@
+pillow
+numpy
+scipy
+matplotlib
+streamlit-drawable-canvas
+accelerate
+torch==2.0.1
+transformers==4.30.2
+diffusers==0.11.1