Update README.md

README.md

@@ -117,8 +117,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 num_images_per_prompt = 1
 print(f"Running on: {device}")
 
-prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to(device)
-decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=torch.float16).to(device)
+prior = StableCascadePriorPipeline.from_pretrained("Ketengan-Diffusion/SomniumSC-v1.1", torch_dtype=torch.bfloat16).to(device) # point to the fine tuned model that you desired (stage C)
+decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=torch.float16).to(device) # point to the "Mother" model which is from stabilityai (Stage A and B)
 
 prompt = "An Astronout riding a horse"
 negative_prompt = ""