Pref-Restoration / DiffusionNFT / model / examples.py
examples.py
Raw
"""
Usage examples: Various calling methods for TextToImagePipeline
"""

from DiffusionNFT.model import TextToImagePipeline, load_pipeline
from PIL import Image

def example_basic_usage():
    """Basic usage example"""
    # Method 1: Direct initialization
    pipeline = TextToImagePipeline.from_pretrained(
        model_path="/data/phd/yaozhengjian/zjYao_Exprs/BLIP-3o-next/Face-Restore_restoration-FFHQ+CelebA/checkpoint-30800",
        device="cuda:0"
    )
    
    # Generate image
    result = pipeline(
        prompt="Please reconstruct the given image.",
        image="path/to/your/image.jpg"
    )
    
    # Save result
    result.save("output.jpg")
    print(f"Generated {len(result.images)} images")


def example_with_pil_image():
    """Example using PIL Image object"""
    # Load pipeline
    pipeline = load_pipeline(
        model_path="your_model_path",
        device="cuda:0",
        dtype="bfloat16"
    )
    
    # Use PIL Image object
    input_image = Image.open("input.jpg")
    result = pipeline(
        prompt="Please enhance and reconstruct this image.",
        image=input_image
    )
    
    # Access original image and generated images
    original = result.original_image
    generated = result.images[0]
    
    # Save
    generated.save("enhanced_output.jpg")


def example_batch_processing():
    """Batch processing example"""
    import os
    from pathlib import Path
    
    # Initialize pipeline
    pipeline = TextToImagePipeline.from_pretrained("your_model_path")
    
    # Batch process images in a folder
    input_folder = Path("input_images")
    output_folder = Path("output_images")
    output_folder.mkdir(exist_ok=True)
    
    for image_file in input_folder.glob("*.jpg"):
        try:
            result = pipeline(
                prompt="Please reconstruct the given image.",
                image=str(image_file)
            )
            
            output_path = output_folder / f"restored_{image_file.name}"
            result.save(str(output_path))
            print(f"Processed: {image_file.name}")
            
        except Exception as e:
            print(f"Error processing {image_file.name}: {e}")


def example_custom_config():
    """Custom configuration example"""
    from DiffusionNFT.model import PipelineConfig
    
    # Create custom configuration
    config = PipelineConfig(
        model_path="your_model_path",
        device="cuda:1",  # Use a different GPU
        scale=1,          # Adjust scale parameter
        seq_len=1024,     # Adjust sequence length
        top_p=0.9,        # Adjust sampling parameters
        top_k=1000,
        image_size=1024   # Use a larger image size
    )
    
    pipeline = TextToImagePipeline(config)
    
    result = pipeline(
        prompt="Please reconstruct and enhance this image with high quality.",
        image="high_res_input.jpg"
    )
    
    result.save("high_quality_output.jpg")


if __name__ == "__main__":
    # Run examples (please modify paths according to actual situation)
    print("TextToImagePipeline usage examples")
    print("Please modify the paths in the code according to your actual model path and image path")
    
    # example_basic_usage()
    # example_with_pil_image() 
    # example_batch_processing()
    # example_custom_config()