Housegan

class HouseGanConfig(PretrainedConfig):
    """Configuration for the House-GAN graph-conditioned generator.

    Args:
        dataset_name: Dataset identifier for the vectorized floorplan assets.
        target_set: House-GAN split target set, one of ``A`` through ``E``.
        checkpoint_step: Original checkpoint training step.
        id2label: Public zero-based room label map.
        relation_id2label: Signed relation label map.
        latent_dim: Per-room latent vector dimension.
        node_feature_dim: One-hot room feature dimension.
        graph_edge_values: Supported signed edge values.
        mask_size: Generated square room-mask size.
        canvas_size: Original floorplan canvas size as ``(width, height)``.
        cmp_channels: CMP feature channel count.
        num_cmp_layers: Number of CMP layers in the generator.
        postprocess_threshold: Threshold used by mask-to-box postprocessing.
        bbox_source: Source of public boxes.
        source_checkpoint: Original checkpoint path or name.
        conversion_report: Measured conversion metadata.
        license_note: Upstream license and research-purpose warning.

    Examples:
        >>> config = HouseGanConfig()
        >>> config.id2label[0]
        'living_room'
    """

    model_type = "housegan"

    def __init__(
        self,
        *,
        dataset_name: str = "housegan_floorplan_vectorized",
        target_set: str = "D",
        checkpoint_step: int = 200000,
        id2label: Id2LabelMapping | None = None,
        relation_id2label: Id2LabelMapping | None = None,
        latent_dim: int = 128,
        node_feature_dim: int = 10,
        graph_edge_values: tuple[int, int] = (-1, 1),
        mask_size: int = 32,
        canvas_size: tuple[int, int] = (256, 256),
        cmp_channels: int = 16,
        num_cmp_layers: int = 2,
        postprocess_threshold: float = 0.0,
        bbox_source: str = "generated_mask",
        source_checkpoint: str | None = None,
        conversion_report: HouseGanConversionReport | None = None,
        license_note: str = "GPL-3.0 with upstream research-purpose notice",
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize a House-GAN config."""
        kwargs.pop("label2id", None)
        kwargs.pop("num_labels", None)
        kwargs.pop("max_supported_room_type_id", None)
        self.dataset_name = dataset_name
        self.target_set = target_set
        self.checkpoint_step = checkpoint_step
        self.id2label = {
            int(key): value for key, value in (id2label or DEFAULT_ID2LABEL).items()
        }
        self.label2id = {value: key for key, value in self.id2label.items()}
        self.relation_id2label = {
            int(key): value
            for key, value in (relation_id2label or DEFAULT_RELATION_ID2LABEL).items()
        }

        self.latent_dim = latent_dim
        self.node_feature_dim = node_feature_dim
        self.graph_edge_values = tuple(graph_edge_values)
        self.mask_size = mask_size
        self.canvas_size = tuple(canvas_size)

        self.cmp_channels = cmp_channels
        self.num_cmp_layers = num_cmp_layers
        self.postprocess_threshold = postprocess_threshold
        self.bbox_source = bbox_source
        self.source_checkpoint = source_checkpoint

        self.conversion_report = conversion_report or {}
        self.license_note = license_note
        self.num_labels = len(self.id2label)
        self.max_supported_room_type_id = self.num_labels - 1

        super().__init__(id2label=self.id2label, label2id=self.label2id)
        for key, value in kwargs.items():
            setattr(self, key, value)

@dataclass(frozen=True)
class HouseGanRelation:
    """Adjacency relation between two room nodes."""

    source: int
    target: int
    adjacent: bool
    weight: float | None = None

@dataclass(frozen=True)
class HouseGanRoomNode:
    """Room node in a House-GAN scene graph."""

    id: int
    label: int | str
    bbox: tuple[float, float, float, float] | None = None
    attributes: Mapping[str, object] | None = None

@dataclass(frozen=True)
class HouseGanSceneGraph:
    """Flat room relation graph used by House-GAN."""

    nodes: tuple[HouseGanRoomNode, ...]
    relations: tuple[HouseGanRelation, ...] | None = None

class HouseGanGenerator(PreTrainedModel):
    """Transformers-compatible House-GAN generator.

    Args:
        config: House-GAN configuration.

    Examples:
        >>> model = HouseGanGenerator(HouseGanConfig())
        >>> latents = torch.zeros(2, 128)
        >>> nodes = torch.eye(10)[:2]
        >>> edges = torch.tensor([[0, 1, 1]])
        >>> tuple(model(latents, nodes, edges).masks.shape)
        (2, 32, 32)
    """

    config_class = HouseGanConfig
    base_model_prefix = "housegan"
    main_input_name = "latents"
    supports_gradient_checkpointing = False

    def __init__(self, config: HouseGanConfig) -> None:
        """Initialize generator layers."""
        super().__init__(config)
        init_size = config.mask_size // 4
        in_features = config.latent_dim + config.node_feature_dim
        channels = config.cmp_channels
        self.init_size = init_size
        self.l1 = nn.Sequential(nn.Linear(in_features, channels * init_size**2))
        self.upsample_1 = nn.Sequential(
            *_conv_block(channels, channels, 4, 2, 1, act="leaky", upsample=True)
        )
        self.upsample_2 = nn.Sequential(
            *_conv_block(channels, channels, 4, 2, 1, act="leaky", upsample=True)
        )
        self.cmp_1 = CMP(channels)
        self.cmp_2 = CMP(channels)
        self.decoder = nn.Sequential(
            *_conv_block(channels, 256, 3, 1, 1, act="leaky"),
            *_conv_block(256, 128, 3, 1, 1, act="leaky"),
            *_conv_block(128, 1, 3, 1, 1, act="tanh"),
        )
        self.post_init()

    def forward(
        self,
        latents: Float[torch.Tensor, "elements latent"],
        node_features: Float[torch.Tensor, "elements room_labels"],
        edges: Int[torch.Tensor, "edges 3"],
        return_dict: bool | None = None,
    ) -> HouseGanModelOutput | tuple[Float[torch.Tensor, "elements height width"]]:
        """Run a House-GAN forward pass.

        Args:
            latents: Per-room latent vectors.
            node_features: Per-room one-hot room features.
            edges: Signed complete graph triples.
            return_dict: Whether to return ``HouseGanModelOutput``.

        Returns:
            Raw generated room masks.
        """
        if latents.ndim != 2 or latents.shape[-1] != self.config.latent_dim:
            raise ValueError("latents must have shape (elements, latent_dim)")

        if node_features.shape != (latents.shape[0], self.config.node_feature_dim):
            raise ValueError(
                "node_features must have shape (elements, node_feature_dim)"
            )

        if edges.ndim != 2 or edges.shape[-1] != 3:
            raise ValueError("edges must have shape (edges, 3)")

        dtype = next(self.parameters()).dtype
        latents = latents.to(dtype=dtype, device=self.device)
        node_features = node_features.to(dtype=dtype, device=self.device)
        edges = edges.to(dtype=torch.long, device=self.device)
        hidden = torch.cat(
            [latents.view(-1, self.config.latent_dim), node_features], dim=1
        )
        hidden = self.l1(hidden)
        hidden = hidden.view(
            -1, self.config.cmp_channels, self.init_size, self.init_size
        )
        hidden = self.cmp_1(hidden, edges).view(-1, *hidden.shape[1:])
        hidden = self.upsample_1(hidden)
        hidden = self.cmp_2(hidden, edges).view(-1, *hidden.shape[1:])
        hidden = self.upsample_2(hidden)
        masks = self.decoder(hidden.view(-1, hidden.shape[1], *hidden.shape[2:]))
        masks = masks.view(-1, *masks.shape[2:])
        if return_dict is False:
            return (masks,)
        return HouseGanModelOutput(
            masks=masks, node_features=node_features, edges=edges
        )

@dataclass
class HouseGanModelOutput(ModelOutput):
    """Raw House-GAN model output."""

    masks: Float[torch.Tensor, "elements height width"]
    node_features: Float[torch.Tensor, "elements room_labels"] | None = None
    edges: Int[torch.Tensor, "edges 3"] | None = None

class HouseGanPipeline(LayoutGenerationPipeline):
    """Transformers-side House-GAN layout generation pipeline."""

    config_class: ClassVar[type[PretrainedConfig]] = HouseGanConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = (
        model_processor_component_specs(
            model_loader=_load_model_component,
            processor_loader=_load_processor_component,
        )
    )

    config: HouseGanConfig
    model: HouseGanGenerator
    processor: HouseGanProcessor

    def __init__(
        self,
        model: HouseGanGenerator,
        processor: HouseGanProcessor | None = None,
        config: HouseGanConfig | None = None,
        device: int | torch.device | None = None,
    ) -> None:
        """Initialize the pipeline."""
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor or HouseGanProcessor(
            config=self.config,
        )
        if device is not None:
            resolved = (
                torch.device("cpu")
                if isinstance(device, int) and device < 0
                else torch.device(f"cuda:{device}")
                if isinstance(device, int)
                else device
            )
            self.to(resolved)

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PipelineComponent | None],
    ) -> "HouseGanPipeline":
        """Build a pipeline from loaded model and processor components."""
        return cls(
            config=cast(HouseGanConfig, config),
            model=cast(HouseGanGenerator, components["model"]),
            processor=cast(HouseGanProcessor, components["processor"]),
        )

    @torch.no_grad()
    def __call__(  # ty: ignore[invalid-method-override]
        self,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.relation,
        labels: Int[torch.Tensor, "..."]
        | Int[np.ndarray, "..."]
        | NestedIntList
        | None = None,
        bbox: Float[torch.Tensor, "... 4"]
        | Float[np.ndarray, "... 4"]
        | NestedFloatList
        | None = None,
        mask: Bool[torch.Tensor, "..."]
        | Bool[np.ndarray, "..."]
        | NestedBoolList
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "..."] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        num_inference_steps: int | None = None,
        scene_graph: HouseGanSceneGraph
        | HouseGanSceneGraphPayload
        | list[HouseGanSceneGraphPayload]
        | None = None,
        relations: Sequence[HouseGanRelationPayload] | None = None,
        latents: Float[torch.Tensor, "elements latent"] | None = None,
        output_type: OutputType = "dataclass",
        return_intermediates: bool = False,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Shaped[torch.Tensor, "..."]
            | dict[int, str]
            | list[dict[str, str | int | float | bool | None]]
            | None,
        ]
    ):
        """Generate a floorplan layout from room relation constraints."""
        del num_elements, num_inference_steps
        if batch_size < 1:
            raise ValueError("batch_size must be positive")

        graph_batch = _expand_graph_batch(scene_graph, batch_size)
        outputs: list[LayoutGenerationOutput] = []
        torch_generator = self.prepare_generator(
            generator=generator,
            seed=seed,
            device=self.device or next(self.model.parameters()).device,
        )
        for graph_index, graph_item in enumerate(graph_batch):
            condition = self.processor(
                condition_type=condition_type,
                scene_graph=graph_item,
                relations=relations,
                labels=labels,
                bbox=bbox,
                mask=mask,
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
            )
            node_features = condition["node_features"].to(self.model.device)
            edges = condition["edges"].to(self.model.device)
            labels_t = condition["labels"].to(self.model.device)
            room_count = node_features.shape[0]
            graph_latents = latents
            if graph_latents is None:
                graph_latents = torch.randn(
                    room_count,
                    self.model.config.latent_dim,
                    generator=torch_generator,
                    device=self.model.device,
                    dtype=next(self.model.parameters()).dtype,
                )
            elif graph_latents.ndim == 3:
                graph_latents = graph_latents[graph_index]
            model_output = self.model(
                latents=graph_latents.to(self.model.device),
                node_features=node_features,
                edges=edges,
            )
            decoded = self.processor.post_process_masks(
                model_output.masks,
                labels=labels_t,
                edges=edges,
                node_features=node_features,
                scene_graph=condition["scene_graph"],
                output_type="dataclass",
                return_intermediates=return_intermediates,
            )
            outputs.append(cast(LayoutGenerationOutput, decoded))
        merged = _merge_outputs(outputs, output_type=output_type)
        return merged

class HouseGanProcessor(ProcessorMixin):
    """Normalize House-GAN scene graphs and decode generated masks."""

    attributes: list[str] = []
    config_name = "processor_config.json"

    def __init__(
        self,
        *,
        config: HouseGanConfig,
        default_missing_relation: Literal["not_adjacent", "error"] = "not_adjacent",
    ) -> None:
        """Initialize processor metadata."""
        self.config = config
        self.id2label = {
            int(key): value
            for key, value in cast(Id2LabelMapping, self.config.id2label).items()
        }
        self.label2id = {value: key for key, value in self.id2label.items()}
        self.relation_id2label = {
            int(key): value
            for key, value in cast(
                Id2LabelMapping, self.config.relation_id2label
            ).items()
        }
        self.canvas_size = tuple(self.config.canvas_size)
        self.mask_size = self.config.mask_size
        self.default_missing_relation = default_missing_relation
        self.chat_template = None

    def save_pretrained(
        self,
        save_directory: str | Path,
        push_to_hub: bool = False,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Save processor metadata."""
        del push_to_hub, kwargs
        root = Path(save_directory)
        root.mkdir(parents=True, exist_ok=True)
        payload = {
            "config": self.config.to_dict(),
            "processor_class": self.__class__.__name__,
            "id2label": self.id2label,
            "relation_id2label": self.relation_id2label,
            "canvas_size": self.canvas_size,
            "mask_size": self.mask_size,
            "default_missing_relation": self.default_missing_relation,
        }
        (root / self.config_name).write_text(
            json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8"
        )

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | PathLike[str],
        cache_dir: str | PathLike[str] | None = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: str | bool | None = None,
        revision: str = "main",
        subfolder: str | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> Self:
        """Load processor metadata from ``processor_config.json``."""
        del cache_dir, force_download, local_files_only, token, revision, kwargs
        root = Path(pretrained_model_name_or_path)
        if subfolder is not None:
            root = root / subfolder
        payload = json.loads((root / cls.config_name).read_text(encoding="utf-8"))
        config_payload = payload.get("config")
        if not isinstance(config_payload, dict):
            raise TypeError("processor config payload must be a dictionary")

        return cls(
            config=HouseGanConfig.from_dict(config_payload),
            default_missing_relation=payload.get(
                "default_missing_relation", "not_adjacent"
            ),
        )

    def __call__(
        self,
        *,
        condition_type: ConditionType | str = ConditionType.relation,
        scene_graph: HouseGanSceneGraph | HouseGanSceneGraphPayload | None = None,
        relations: Sequence[HouseGanRelationPayload] | None = None,
        labels: Int[torch.Tensor, "..."]
        | Int[np.ndarray, "..."]
        | Sequence[Sequence[int]]
        | Sequence[int]
        | None = None,
        bbox: Float[torch.Tensor, "... 4"]
        | Float[np.ndarray, "... 4"]
        | Sequence[Sequence[Sequence[float]]]
        | Sequence[Sequence[float]]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "..."]
        | Bool[np.ndarray, "..."]
        | Sequence[Sequence[bool]]
        | Sequence[bool]
        | Sequence[ArrayLikeInput]
        | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Encode public relation inputs into House-GAN tensors."""
        if return_tensors != "pt":
            raise ValueError("HouseGanProcessor only supports return_tensors='pt'")

        self.normalize_condition_type(condition_type)
        relation_payload = relations
        if relation_payload is None and bbox is not None and labels is not None:
            bbox_t, labels_t, _ = prepare_layout_tensors(
                bbox=cast(
                    Float[torch.Tensor, "... 4"]
                    | Float[np.ndarray, "... 4"]
                    | Sequence[Sequence[Sequence[float]]]
                    | Sequence[Sequence[float]]
                    | Sequence[ArrayLikeInput],
                    bbox,
                ),
                labels=cast(
                    Int[torch.Tensor, "..."]
                    | Int[np.ndarray, "..."]
                    | Sequence[Sequence[int]]
                    | Sequence[int]
                    | Sequence[ArrayLikeInput],
                    labels,
                ),
                mask=cast(
                    Bool[torch.Tensor, "..."]
                    | Bool[np.ndarray, "..."]
                    | Sequence[Sequence[bool]]
                    | Sequence[bool]
                    | Sequence[ArrayLikeInput]
                    | None,
                    mask,
                ),
                box_format=normalize_box_format(box_format),
                normalized=normalized,
                canvas_size=canvas_size or self.canvas_size,
                clamp_converted_normalized=True,
            )
            relation_payload = relation_from_bboxes(
                _xywh_to_ltrb_list(bbox_t[0]),
            )
            labels = labels_t[0]
        graph = normalize_scene_graph(
            scene_graph,
            labels=labels,
            relations=relation_payload,
            id2label=self.id2label,
        )
        if self.default_missing_relation == "error" and not graph.relations:
            raise ValueError(
                "House-GAN requires relations when missing-pair policy is 'error'"
            )

        node_features = graph_to_node_features(
            graph.nodes,
            label2id=self.label2id,
            num_labels=len(self.id2label),
        )
        edges = complete_signed_edges(
            graph.nodes,
            graph.relations,
            default_adjacent=False,
        )
        label_ids = torch.tensor(
            [
                self.label2id[node.label]
                if isinstance(node.label, str)
                else int(node.label)
                for node in graph.nodes
            ],
            dtype=torch.long,
        )
        return BatchEncoding(
            {
                "node_features": node_features,
                "edges": edges,
                "labels": label_ids,
                "scene_graph": graph,
            }
        )

    def normalize_condition_type(
        self, condition_type: ConditionType | str
    ) -> ConditionType:
        """Normalize and validate the House-GAN condition type."""
        condition = normalize_condition_type(condition_type)
        if condition is not ConditionType.relation:
            raise NotImplementedError(
                "House-GAN supports only condition_type='relation' and aliases "
                "'scene_graph', 'graph', or 'gen_r'."
            )

        return condition

    def post_process_masks(
        self,
        masks: Float[torch.Tensor, "elements height width"],
        *,
        labels: Int[torch.Tensor, "elements"],
        edges: Int[torch.Tensor, "edges 3"] | None = None,
        node_features: Float[torch.Tensor, "elements room_labels"] | None = None,
        scene_graph: HouseGanSceneGraph | None = None,
        output_type: OutputType = "dataclass",
        return_intermediates: bool = False,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Shaped[torch.Tensor, "..."]
            | dict[int, str]
            | dict[str, Shaped[torch.Tensor, "..."] | HouseGanSceneGraph | None]
            | None,
        ]
    ):
        """Convert generated room masks to public normalized boxes."""
        bbox_ltrb = mask_to_ltrb(masks, threshold=0.0)
        bbox_xywh = (
            ltrb_to_xywh(bbox_ltrb / float(self.mask_size)).unsqueeze(0).clamp(0.0, 1.0)
        )
        labels_b = labels.to(dtype=torch.long).unsqueeze(0)
        valid = torch.ones(labels_b.shape, dtype=torch.bool, device=labels_b.device)
        intermediates = None
        if return_intermediates:
            intermediates = {
                "room_masks": masks.detach(),
                "bbox_ltrb_32": bbox_ltrb,
                "signed_edges": edges,
                "node_features": node_features,
                "scene_graph": scene_graph,
            }
        if output_type == "dict":
            return {
                "bbox": bbox_xywh,
                "labels": labels_b,
                "mask": valid,
                "id2label": self.id2label,
                "sequences": edges,
                "scores": None,
                "trajectory": None,
                "intermediates": intermediates,
            }
        return LayoutGenerationOutput(
            bbox=bbox_xywh,
            labels=labels_b,
            mask=valid,
            id2label=self.id2label,
            sequences=edges,
            intermediates=intermediates,
        )