zurizheng/scribblediffusion-fruit-test

ScribbleDiffusion: Fruit Dataset Fine-tuned Model

This model is a fine-tuned version of Stable Diffusion v1.5 for sketch-to-image generation, specifically trained on a fruit dataset.

Model Description

ScribbleDiffusion allows you to generate high-quality images from simple sketches combined with text prompts. This version has been trained on 6 fruit categories:

• Apple
• Banana
• Guava
• Lime
• Orange
• Pomegranate

Model Architecture

The model consists of three main components:

1. UNet: Modified Stable Diffusion 1.5 UNet for diffusion generation
2. Sketch Encoder: Cross-attention encoder that processes input sketches
3. Sketch Text Combiner: Module that combines sketch and text embeddings

Training Details

• Base Model: Stable Diffusion v1.5
• Training Steps: 5,000 steps
• Batch Size: 1 (with gradient accumulation steps: 8)
• Learning Rate: 0.0001
• Image Resolution: 256x256
• Mixed Precision: FP16
• Optimizer: AdamW
• Scheduler: Cosine with 500 warmup steps

Usage

Loading the Model

import torch
from diffusers import UNet2DConditionModel, AutoencoderKL, DDIMScheduler
from transformers import CLIPTokenizer, CLIPTextModel
from safetensors.torch import load_file

# Load base models
vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
scheduler = DDIMScheduler.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="scheduler")

# Load fine-tuned UNet
unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
unet_state = load_file("unet.safetensors")
unet.load_state_dict(unet_state)

# Load custom sketch components
# Note: You'll need the custom SketchCrossAttentionEncoder and SketchTextCombiner classes
sketch_encoder_state = load_file("sketch_encoder.safetensors")
sketch_combiner_state = load_file("sketch_text_combiner.safetensors")

Inference Example

# Your sketch should be a 1-channel grayscale image (edges/contours)
sketch = load_sketch_image("your_sketch.png")  # 256x256 grayscale
prompt = "a red apple"

# Process sketch and text
sketch_embeddings = sketch_encoder(sketch)
text_embeddings = text_encoder(tokenize(prompt))
combined_embeddings = sketch_text_combiner(text_embeddings, sketch_embeddings)

# Generate image
with torch.no_grad():
    latents = torch.randn((1, 4, 32, 32))  # 256x256 -> 32x32 latents
    for t in scheduler.timesteps:
        noise_pred = unet(latents, t, encoder_hidden_states=combined_embeddings).sample
        latents = scheduler.step(noise_pred, t, latents).prev_sample
    
    # Decode to image
    image = vae.decode(latents / vae.config.scaling_factor).sample

Model Files

• unet.safetensors (3.3GB): Fine-tuned UNet model weights
• sketch_encoder.safetensors (24MB): Sketch encoder weights
• sketch_text_combiner.safetensors (16 bytes): Sketch-text combiner weights
• training_info.json: Training metadata

Training Data

The model was trained on a curated fruit dataset containing high-quality images of 6 fruit categories. Sketches were automatically generated using edge detection algorithms.

Limitations

• Trained specifically on fruit images (may not generalize well to other objects)
• Input resolution limited to 256x256
• Requires specific sketch preprocessing (edge detection style)
• Best results with simple, clear sketches

Ethics and Bias

This model inherits potential biases from the Stable Diffusion v1.5 base model. The fruit dataset training may introduce biases toward specific fruit appearances and orientations present in the training data.

Citation

@misc{scribblediffusion-fruit-2024,
  title={ScribbleDiffusion: Fruit Dataset Fine-tuned Model},
  author={Your Name},
  year={2024},
  howpublished={\\url{https://huggingface.co/your-username/scribblediffusion-fruit}}
}

License

This model is released under the same license as Stable Diffusion v1.5. Please refer to the original licensing terms.

Acknowledgments

• Based on Stable Diffusion v1.5 by Runway ML
• Training infrastructure and optimization techniques
• Fruit dataset compilation and preprocessing