Update README.md

README.md

@@ -22,34 +22,46 @@ inference: true
 should probably proofread and complete it, then remove this comment. -->
 
 
-# controlnet-sidnarsipur/controlnet_rough
+# controlnet_rough
 
-These are controlnet weights trained on stabilityai/stable-diffusion-2-1-base with new type of conditioning.
-You can find some example images below.
+Generate a roughness map from a photograph or basecolor (albedo) map.
 
-prompt: Roughness Map
-![images_0)](./images_0.png)
-prompt: Roughness Map
-![images_1)](./images_1.png)
-prompt: Roughness Map
-![images_2)](./images_2.png)
-prompt: Roughness Map
-![images_3)](./images_3.png)
+# Usage
 
+```
+import argparse
 
+from PIL import Image
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from diffusers.utils import load_image
+import torch
 
-## Intended uses & limitations
+parser = argparse.ArgumentParser(description="Args for parser")
+parser.add_argument("--seed", type=int, default=1, help="Seed for inference")
+args = parser.parse_args()
 
-#### How to use
+base_model_path = "stabilityai/stable-diffusion-2-1-base"
+controlnet_path = "sidnarsipur/controlnet_rough"
 
-```python
-# TODO: add an example code snippet for running this diffusion pipeline
-```
+controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    base_model_path, controlnet=controlnet, torch_dtype=torch.float16
+)
 
-#### Limitations and bias
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.enable_xformers_memory_efficient_attention()
+pipe.enable_model_cpu_offload()
 
-[TODO: provide examples of latent issues and potential remediations]
+control_image = load_image("inference/basecolor.png") #Change based on your image path
+prompt = "Roughness Map" #Don't change!
 
-## Training details
+if control_image.size[0] > 2048 or control_image.size[1] > 2048: #Optional
+    control_image = control_image.resize((control_image.size[0] // 2, control_image.size[1] // 2))
 
-[TODO: describe the data used to train the model]
+generator = torch.manual_seed(args.seed)
+
+image = pipe(
+    prompt, num_inference_steps=50, generator=generator, image=control_image
+).images[0]
+image.save("inference/normal.png")
+```