Update README.md

README.md

@@ -31,23 +31,24 @@ A LoRA fine-tuned [Stable Diffusion XL (SDXL)](https://huggingface.co/stabilitya
 ## 📦 How to Use
 
 ```python
-from diffusers import StableDiffusionXLPipeline
+pip install diffusers
+
+```python
+from diffusers import DiffusionPipeline
 import torch
 
 # Load base SDXL model
-pipeline = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-).to("cuda")
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+pipe.to("cuda")
 
 # Load LoRA weights
-pipeline.load_lora_weights("aryanbaghel/ui-generator-prototype")
+pipe.load_lora_weights("aryanbaghel/ui-generator-prototype")
 
 # Enable LoRA
-pipeline.fuse_lora()
+pipe.fuse_lora()
 
 # Generate a UI screen
 prompt = "A mobile UI showing a hotel reservation screen with elegant typography and clean layout"
-image = pipeline(prompt=prompt).images[0]
-image.save("output.png")
+image = pipe(prompt=prompt).images[0]
+image.save("output.png")
+print("SUCCESS!")