Update README.md

README.md

@@ -21,4 +21,118 @@ image = pipe(
     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")
+```