Update README.md

README.md

@@ -33,7 +33,21 @@ prompt: An image of breast cancer histopathology with detailed cellular structur
 #### How to use
 
 ```python
-# TODO: add an example code snippet for running this diffusion pipeline
+controlnet = ControlNetModel.from_pretrained("RiddleHe/SD14_pathology_controlnet", torch_dtype=torch.float16)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4", controlnet=controlnet, torch_dtype=torch.float16
+)
+pipe.load_lora_weights("RiddleHe/SD14_pathology_lora")  # Load lora weights for pathology image generation
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.to('cuda')
+
+prompt = "A histopathology image of breast cancer tissue."
+mask = mask.convert("RGB")  # Provide a mask
+
+generator = torch.Generator(device='cuda').manual_seed(42)
+
+with torch.no_grad():
+  out = pipe(prompt, image=mask, num_inference_steps=70, num_images_per_prompt=3, generator=generator).images
 ```
 
 #### Limitations and bias
@@ -42,4 +56,6 @@ prompt: An image of breast cancer histopathology with detailed cellular structur
 
 ## Training details
 
-[TODO: describe the data used to train the model]
+The model is trained on 28216 image-mask pairs from the BRCA breast cancer dataset. Input is mask and output is image. 
+
+Mask is a single channel image with integer values from 0 to 21 representing 22 classes, eg. 1 representing tumor, 2 representing stroma.