You can deploy this model and start earning money today!

zurizheng/
scribblediffusion-fruit

Name: scribblediffusion-fruit
Availability: InStock

scribblediffusion

Views

No views yet

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:

UNet: Modified Stable Diffusion 1.5 UNet for diffusion generation
Sketch Encoder: Cross-attention encoder that processes input sketches
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

python

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

python

# 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

bibtex

@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

zurizheng/scribblediffusion-fruit