| --- |
| library_name: diffusers |
| --- |
| |
| # yujiepan/stable-diffusion-3-tiny-random |
|
|
| This pipeline is intended for debugging. It is adapted from [stabilityai/stable-diffusion-3-medium-diffusers](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) with smaller size and randomly initialized parameters. |
|
|
| ## Usage |
| ```python |
| import torch |
| from diffusers import StableDiffusion3Pipeline |
| |
| pipe = StableDiffusion3Pipeline.from_pretrained("yujiepan/stable-diffusion-3-tiny-random", torch_dtype=torch.float16) |
| pipe = pipe.to("cuda") |
| |
| image = pipe( |
| "A cat holding a sign that says hello world", |
| negative_prompt="", |
| num_inference_steps=2, |
| guidance_scale=7.0, |
| ).images[0] |
| image |
| ``` |
|
|
| ## Codes |
| ```python |
| import importlib |
| |
| import torch |
| import transformers |
| |
| import diffusers |
| import rich |
| |
| |
| def get_original_model_configs(pipeline_cls: type[diffusers.DiffusionPipeline], pipeline_id: str): |
| pipeline_config: dict[str, list[str]] = pipeline_cls.load_config(pipeline_id) |
| model_configs = {} |
| |
| for subfolder, import_strings in pipeline_config.items(): |
| if subfolder.startswith("_"): |
| continue |
| module = importlib.import_module(".".join(import_strings[:-1])) |
| cls = getattr(module, import_strings[-1]) |
| if issubclass(cls, transformers.PreTrainedModel): |
| config_class: transformers.PretrainedConfig = cls.config_class |
| config = config_class.from_pretrained(pipeline_id, subfolder=subfolder) |
| model_configs[subfolder] = config |
| elif issubclass(cls, diffusers.ModelMixin) and issubclass(cls, diffusers.ConfigMixin): |
| config = cls.load_config(pipeline_id, subfolder=subfolder) |
| model_configs[subfolder] = config |
| |
| return model_configs |
| |
| |
| def load_pipeline(pipeline_cls: type[diffusers.DiffusionPipeline], pipeline_id: str, model_configs: dict[str, dict]): |
| pipeline_config: dict[str, list[str]] = pipeline_cls.load_config(pipeline_id) |
| components = {} |
| for subfolder, import_strings in pipeline_config.items(): |
| if subfolder.startswith("_"): |
| continue |
| module = importlib.import_module(".".join(import_strings[:-1])) |
| cls = getattr(module, import_strings[-1]) |
| print(f"Loading:", ".".join(import_strings)) |
| if issubclass(cls, transformers.PreTrainedModel): |
| config = model_configs[subfolder] |
| component = cls(config) |
| elif issubclass(cls, transformers.PreTrainedTokenizerBase): |
| component = cls.from_pretrained(pipeline_id, subfolder=subfolder) |
| elif issubclass(cls, diffusers.ModelMixin) and issubclass(cls, diffusers.ConfigMixin): |
| config = model_configs[subfolder] |
| component = cls.from_config(config) |
| elif issubclass(cls, diffusers.SchedulerMixin) and issubclass(cls, diffusers.ConfigMixin): |
| component = cls.from_pretrained(pipeline_id, subfolder=subfolder) |
| else: |
| raise (f"unknown {subfolder}: {import_strings}") |
| components[subfolder] = component |
| pipeline = pipeline_cls(**components) |
| return pipeline |
| |
| |
| def get_pipeline(): |
| torch.manual_seed(42) |
| pipeline_id = "stabilityai/stable-diffusion-3-medium-diffusers" |
| pipeline_cls = diffusers.StableDiffusion3Pipeline |
| model_configs = get_original_model_configs(pipeline_cls, pipeline_id) |
| rich.print(model_configs) |
| |
| HIDDEN_SIZE = 8 |
| |
| model_configs["text_encoder"].hidden_size = HIDDEN_SIZE |
| model_configs["text_encoder"].intermediate_size = HIDDEN_SIZE * 2 |
| model_configs["text_encoder"].num_attention_heads = 2 |
| model_configs["text_encoder"].num_hidden_layers = 2 |
| model_configs["text_encoder"].projection_dim = HIDDEN_SIZE |
| |
| model_configs["text_encoder_2"].hidden_size = HIDDEN_SIZE |
| model_configs["text_encoder_2"].intermediate_size = HIDDEN_SIZE * 2 |
| model_configs["text_encoder_2"].num_attention_heads = 2 |
| model_configs["text_encoder_2"].num_hidden_layers = 2 |
| model_configs["text_encoder_2"].projection_dim = HIDDEN_SIZE |
| |
| model_configs["text_encoder_3"].d_model = HIDDEN_SIZE |
| model_configs["text_encoder_3"].d_ff = HIDDEN_SIZE * 2 |
| model_configs["text_encoder_3"].d_kv = HIDDEN_SIZE // 2 |
| model_configs["text_encoder_3"].num_heads = 2 |
| model_configs["text_encoder_3"].num_layers = 2 |
| |
| model_configs["transformer"]["num_layers"] = 2 |
| model_configs["transformer"]["num_attention_heads"] = 2 |
| model_configs["transformer"]["attention_head_dim"] = HIDDEN_SIZE // 2 |
| model_configs["transformer"]["pooled_projection_dim"] = HIDDEN_SIZE * 2 |
| model_configs["transformer"]["joint_attention_dim"] = HIDDEN_SIZE |
| model_configs["transformer"]["caption_projection_dim"] = HIDDEN_SIZE |
| |
| model_configs["vae"]["layers_per_block"] = 1 |
| model_configs["vae"]["block_out_channels"] = [HIDDEN_SIZE] * 4 |
| model_configs["vae"]["norm_num_groups"] = 2 |
| model_configs["vae"]["latent_channels"] = 16 |
| |
| pipeline = load_pipeline(pipeline_cls, pipeline_id, model_configs) |
| return pipeline |
| |
| |
| pipeline = get_pipeline() |
| image = pipeline( |
| "hello world", |
| negative_prompt="runtime error", |
| num_inference_steps=2, |
| guidance_scale=7.0, |
| ).images[0] |
| |
| |
| pipeline = pipeline.to(torch.float16) |
| pipeline.save_pretrained("/tmp/stable-diffusion-3-tiny-random") |
| pipeline.push_to_hub("yujiepan/stable-diffusion-3-tiny-random") |
| ``` |
|
|