Update README.md

README.md

@@ -126,26 +126,86 @@ You may reuse the base model text encoder for inference.
 
 
 ```python
+import argparse
 import torch
-from diffusers import DiffusionPipeline
+from helpers.models.flux.pipeline import FluxPipeline as DiffusionPipeline
 from lycoris import create_lycoris_from_weights
-
-model_id = 'black-forest-labs/FLUX.1-dev'
-adapter_id = 'pytorch_lora_weights.safetensors' # you will have to download this manually
-lora_scale = 1.0
-wrapper, _ = create_lycoris_from_weights(lora_scale, adapter_id, pipeline.transformer)
-wrapper.merge_to()
-
-prompt = "A photo-realistic image of a cat"
-
-pipeline.to('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
-image = pipeline(
-    prompt=prompt,
-    num_inference_steps=20,
-    generator=torch.Generator(device='cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu').manual_seed(1641421826),
-    width=1776,
-    height=512,
-    guidance_scale=3.0,
-).images[0]
-image.save("output.png", format="PNG")
+from huggingface_hub import hf_hub_download
+
+def generate_image(pipeline, prompt, output_file, num_inference_steps, width, height, guidance_scale, seed, device):
+    # Set device
+    pipeline.to(device)
+
+    # Generate image
+    generator = torch.Generator(device=device).manual_seed(seed)
+    image = pipeline(
+        prompt=prompt,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+    ).images[0]
+
+    # Save image
+    output_file = "output.png"
+    image.save(output_file, format="PNG")
+    print(f"Image saved as {output_file}")
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate images using a custom diffusion pipeline with LoRA weights.")
+    parser.add_argument("--model_id", type=str, default='black-forest-labs/FLUX.1-dev', help="Model ID from Hugging Face Hub.")
+    parser.add_argument("--adapter_id", type=str, default='pytorch_lora_weights.safetensors', help="LoRA weights file.")
+    parser.add_argument("--lora_scale", type=float, default=1.0, help="Scale for LoRA weights.")
+    parser.add_argument("--output_file", type=str, default="output.png", help="Output file name for the generated image.")
+    parser.add_argument("--num_inference_steps", type=int, default=30, help="Number of inference steps.")
+    parser.add_argument("--guidance_scale", type=float, default=3.5, help="Guidance scale for the generation.")
+    parser.add_argument("--seed", type=int, default=1641421826, help="Random seed for reproducibility.")
+    parser.add_argument("--device", type=str, default='cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu', help="Device to run the model on.")
+    
+    args = parser.parse_args()
+
+    # Load model and weights
+    hf_hub_download(repo_id="terminusresearch/flux-lokr-garfield-nomask", filename=args.adapter_id, local_dir="./")
+    pipeline = DiffusionPipeline.from_pretrained(args.model_id, torch_dtype=torch.bfloat16)
+
+    # Apply LoRA weights
+    wrapper, _ = create_lycoris_from_weights(args.lora_scale, args.adapter_id, pipeline.transformer)
+    wrapper.merge_to()
+
+    print("Model loaded successfully. Ready to generate images.")
+    
+    while True:
+        user_input = input("Enter a prompt or 'quit' to exit: ")
+        if user_input.lower() == 'quit':
+            break
+
+        # Check for resolution command
+        if user_input.startswith("resolution:"):
+            resolution = user_input.split(":")[1]
+            width, height = map(int, resolution.split("x"))
+            print(f"Resolution set to {width}x{height}")
+            continue
+
+        prompt = user_input
+        output_file = args.output_file.replace(".png", f"_{prompt.replace(' ', '_')}.png")
+
+        # Use default or previously set resolution
+        width = locals().get('width', 1024)
+        height = locals().get('height', 1024)
+
+        generate_image(
+            pipeline=pipeline,
+            prompt=prompt,
+            output_file=output_file,
+            num_inference_steps=args.num_inference_steps,
+            width=width,
+            height=height,
+            guidance_scale=args.guidance_scale,
+            seed=args.seed,
+            device=args.device
+        )
+
+if __name__ == "__main__":
+    main()
 ```