diff --git a/.gitignore b/.gitignore
index 39b6783..8a7c6bf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -34,13 +34,7 @@ captures/
 
 # IntelliJ
 *.iml
-.idea/workspace.xml
-.idea/tasks.xml
-.idea/gradle.xml
-.idea/assetWizardSettings.xml
-.idea/dictionaries
-.idea/libraries
-.idea/caches
+.idea/*
 
 # Keystore files
 # Uncomment the following line if you do not want to check your keystore files in.
diff --git a/README.md b/README.md
index d5b6875..08dbf1c 100644
--- a/README.md
+++ b/README.md
@@ -75,10 +75,11 @@ stateLayout.content()
 stateLayout.loadingWithContent()
         
 // error/info
-stateLayout.infoImage(R.drawable.ic_android_black_64dp)
-                .infoTitle("Ooops.... :(")
-                .infoMessage("Unexpected error occurred. Please refresh the page!")
-                .infoButton("Refresh", onStateLayoutListener)
+stateLayout.InfoLayoutBuilder()
+                    .infoImage(R.drawable.ic_android_black_64dp)
+                    .infoTitle("Ooops.... :(")
+                    .infoMessage("Unexpected error occurred. Please refresh the page!")
+                    .infoButton("Refresh", onStateLayoutListener)
 // error/info 
 stateLayout.info()
 ``` 
diff --git a/library/src/main/java/com/erkutaras/statelayout/StateLayout.kt b/library/src/main/java/com/erkutaras/statelayout/StateLayout.kt
index e63e508..8018100 100644
--- a/library/src/main/java/com/erkutaras/statelayout/StateLayout.kt
+++ b/library/src/main/java/com/erkutaras/statelayout/StateLayout.kt
@@ -1,10 +1,9 @@
 package com.erkutaras.statelayout
 
 import android.content.Context
-import android.support.annotation.LayoutRes
 import android.util.AttributeSet
-import android.view.LayoutInflater
 import android.view.View
+import android.view.ViewStub
 import android.widget.Button
 import android.widget.FrameLayout
 import android.widget.ImageView
@@ -15,13 +14,12 @@ import android.widget.TextView
  */
 class StateLayout @JvmOverloads constructor(context: Context,
                                             attrs: AttributeSet? = null,
-                                            defStyleAttr: Int = 0)
-    : FrameLayout(context, attrs, defStyleAttr) {
+                                            defStyleAttr: Int = 0) : FrameLayout(context, attrs, defStyleAttr) {
 
     private var contentLayout: View? = null
-    private var loadingLayout: View? = null
-    private var infoLayout: View? = null
-    private var loadingWithContentLayout: View? = null
+    private var loadingLayout: ViewStub? = null
+    private var infoLayout: ViewStub? = null
+    private var loadingWithContentLayout: ViewStub? = null
 
     private var state: State = State.CONTENT
 
@@ -41,20 +39,17 @@ class StateLayout @JvmOverloads constructor(context: Context,
     }
 
     private fun setupLoadingState() {
-        loadingLayout = inflate(R.layout.layout_state_loading)
-        loadingLayout?.visibility = View.GONE
+        loadingLayout = ViewStub(context, R.layout.layout_state_loading)
         addView(loadingLayout)
     }
 
     private fun setupInfoState() {
-        infoLayout = inflate(R.layout.layout_state_info)
-        infoLayout?.visibility = View.GONE
+        infoLayout = ViewStub(context, R.layout.layout_state_info)
         addView(infoLayout)
     }
 
     private fun setupLoadingWithContentState() {
-        loadingWithContentLayout = inflate(R.layout.layout_state_loading_with_content)
-        loadingWithContentLayout?.visibility = View.GONE
+        loadingWithContentLayout = ViewStub(context, R.layout.layout_state_loading_with_content)
         addView(loadingWithContentLayout)
     }
 
@@ -76,54 +71,6 @@ class StateLayout @JvmOverloads constructor(context: Context,
         return this
     }
 
-    fun infoImage(imageRes: Int): StateLayout {
-        infoLayout?.findViewById<ImageView>(R.id.imageView_state_layout_info)?.let {
-            it.setImageResource(imageRes)
-            it.visibility = View.VISIBLE
-        }
-        return info()
-    }
-
-    fun infoTitle(title: String): StateLayout {
-        infoLayout?.findViewById<TextView>(R.id.textView_state_layout_info_title)?.let {
-            it.text = title
-            it.visibility = View.VISIBLE
-        }
-        return info()
-    }
-
-    fun infoMessage(message: String): StateLayout {
-        infoLayout?.findViewById<TextView>(R.id.textView_state_layout_info_message)?.let {
-            it.text = message
-            it.visibility = View.VISIBLE
-        }
-        return info()
-    }
-
-    fun infoButtonListener(onStateLayoutListener: OnStateLayoutListener?): StateLayout {
-        infoLayout?.findViewById<Button>(R.id.button_state_layout_info)?.setOnClickListener {
-            onStateLayoutListener?.onStateLayoutInfoButtonClick()
-        }
-        return info()
-    }
-
-    fun infoButtonText(buttonText: String): StateLayout {
-        infoLayout?.findViewById<Button>(R.id.button_state_layout_info)?.let {
-            it.text = buttonText
-            it.visibility = View.VISIBLE
-        }
-        return info()
-    }
-
-    fun infoButton(buttonText: String, onStateLayoutListener: OnStateLayoutListener?): StateLayout {
-        infoLayout?.findViewById<Button>(R.id.button_state_layout_info)?.let { it ->
-            it.text = buttonText
-            it.setOnClickListener { onStateLayoutListener?.onStateLayoutInfoButtonClick() }
-            it.visibility = View.VISIBLE
-        }
-        return info()
-    }
-
     fun info(): StateLayout {
         state = State.INFO
         loadingLayout?.visibility = View.GONE
@@ -148,11 +95,60 @@ class StateLayout @JvmOverloads constructor(context: Context,
         }
     }
 
-    private fun throwChildCountException(): Nothing =
-            throw IllegalStateException("StateLayout can host only one direct child")
+    private fun throwChildCountException(): Nothing = throw IllegalStateException("StateLayout can host only one direct child")
+
+    inner class InfoLayoutBuilder {
+        init {
+            info()
+        }
+
+        fun infoImage(imageRes: Int): InfoLayoutBuilder {
+            findViewById<ImageView>(R.id.imageView_state_layout_info)?.let {
+                it.setImageResource(imageRes)
+                it.visibility = View.VISIBLE
+            }
+            return this
+        }
+
+        fun infoTitle(title: String): InfoLayoutBuilder {
+            infoLayout?.findViewById<TextView>(R.id.textView_state_layout_info_title)?.let {
+                it.text = title
+                it.visibility = View.VISIBLE
+            }
+            return this
+        }
+
+        fun infoMessage(message: String): InfoLayoutBuilder {
+            findViewById<TextView>(R.id.textView_state_layout_info_message)?.let {
+                it.text = message
+                it.visibility = View.VISIBLE
+            }
+            return this
+        }
+
+        fun infoButtonListener(onStateLayoutListener: OnStateLayoutListener?): InfoLayoutBuilder {
+            findViewById<Button>(R.id.button_state_layout_info)?.setOnClickListener {
+                onStateLayoutListener?.onStateLayoutInfoButtonClick()
+            }
+            return this
+        }
+
+        fun infoButtonText(buttonText: String): InfoLayoutBuilder {
+            findViewById<Button>(R.id.button_state_layout_info)?.let {
+                it.text = buttonText
+                it.visibility = View.VISIBLE
+            }
+            return this
+        }
 
-    private fun inflate(@LayoutRes layoutId: Int): View? {
-        return LayoutInflater.from(context).inflate(layoutId, null)
+        fun infoButton(buttonText: String, onStateLayoutListener: OnStateLayoutListener?): InfoLayoutBuilder {
+            findViewById<Button>(R.id.button_state_layout_info)?.let { it ->
+                it.text = buttonText
+                it.setOnClickListener { onStateLayoutListener?.onStateLayoutInfoButtonClick() }
+                it.visibility = View.VISIBLE
+            }
+            return this
+        }
     }
 
     interface OnStateLayoutListener {
diff --git a/library/src/main/res/layout/layout_state_info.xml b/library/src/main/res/layout/layout_state_info.xml
index 1935321..c57a3e0 100644
--- a/library/src/main/res/layout/layout_state_info.xml
+++ b/library/src/main/res/layout/layout_state_info.xml
@@ -1,6 +1,5 @@
 <?xml version="1.0" encoding="utf-8"?>
-<LinearLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:tools="http://schemas.android.com/tools"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
@@ -16,7 +15,7 @@
         android:layout_marginBottom="16dp"
         android:visibility="gone"
         tools:srcCompat="@drawable/ic_info_outline_black_64dp"
-        tools:visibility="visible"/>
+        tools:visibility="visible" />
 
     <TextView
         android:id="@+id/textView_state_layout_info_title"
@@ -28,7 +27,7 @@
         android:textStyle="normal"
         android:visibility="gone"
         tools:text="Oops..."
-        tools:visibility="visible"/>
+        tools:visibility="visible" />
 
     <TextView
         android:id="@+id/textView_state_layout_info_message"
@@ -40,7 +39,7 @@
         android:textStyle="normal"
         android:visibility="gone"
         tools:text="Sorry!... Something goes wrong :("
-        tools:visibility="visible"/>
+        tools:visibility="visible" />
 
     <Button
         android:id="@+id/button_state_layout_info"
@@ -52,5 +51,5 @@
         android:textColor="@android:color/black"
         android:visibility="gone"
         tools:text="Try Again"
-        tools:visibility="visible"/>
+        tools:visibility="visible" />
 </LinearLayout>
diff --git a/library/src/main/res/layout/layout_state_loading.xml b/library/src/main/res/layout/layout_state_loading.xml
index 3d310e6..9345024 100644
--- a/library/src/main/res/layout/layout_state_loading.xml
+++ b/library/src/main/res/layout/layout_state_loading.xml
@@ -1,16 +1,15 @@
 <?xml version="1.0" encoding="utf-8"?>
-<RelativeLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:app="http://schemas.android.com/apk/res-auto"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
-    android:background="?android:colorBackground">
+    android:background="?android:colorBackground"
+    android:gravity="center">
 
     <ProgressBar
         android:id="@+id/progressBar_state_layout"
         android:layout_width="wrap_content"
         android:layout_height="wrap_content"
-        android:layout_centerInParent="true"
-        app:tint="?colorAccent"/>
+        app:tint="?colorAccent" />
 
-</RelativeLayout>
\ No newline at end of file
+</LinearLayout>
\ No newline at end of file
diff --git a/library/src/main/res/layout/layout_state_loading_with_content.xml b/library/src/main/res/layout/layout_state_loading_with_content.xml
index 79d9a12..ce032fc 100644
--- a/library/src/main/res/layout/layout_state_loading_with_content.xml
+++ b/library/src/main/res/layout/layout_state_loading_with_content.xml
@@ -1,6 +1,5 @@
 <?xml version="1.0" encoding="utf-8"?>
-<RelativeLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
+<FrameLayout xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:app="http://schemas.android.com/apk/res-auto"
     android:layout_width="match_parent"
     android:layout_height="match_parent">
@@ -10,8 +9,7 @@
         style="?android:attr/progressBarStyleHorizontal"
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
-        android:layout_alignParentTop="true"
         android:indeterminate="true"
-        app:tint="?colorAccent"/>
+        app:tint="?colorAccent" />
 
-</RelativeLayout>
\ No newline at end of file
+</FrameLayout>
\ No newline at end of file
diff --git a/sample/src/main/java/com/erkutaras/statelayout/sample/StateLayoutSampleActivity.kt b/sample/src/main/java/com/erkutaras/statelayout/sample/StateLayoutSampleActivity.kt
index 63bb24b..f803ac6 100644
--- a/sample/src/main/java/com/erkutaras/statelayout/sample/StateLayoutSampleActivity.kt
+++ b/sample/src/main/java/com/erkutaras/statelayout/sample/StateLayoutSampleActivity.kt
@@ -9,7 +9,8 @@ import android.webkit.WebView
 import android.webkit.WebViewClient
 import android.widget.Toast
 import com.erkutaras.statelayout.StateLayout
-import kotlinx.android.synthetic.main.activity_state_layout_sample.*
+import kotlinx.android.synthetic.main.activity_state_layout_sample.stateLayout
+import kotlinx.android.synthetic.main.activity_state_layout_sample.webView
 
 const val WEB_URL = "http://www.erkutaras.com/"
 
@@ -54,7 +55,8 @@ class StateLayoutSampleActivity : AppCompatActivity(), StateLayout.OnStateLayout
         override fun onReceivedError(view: WebView?, request: WebResourceRequest?, error: WebResourceError?) {
             super.onReceivedError(view, request, error)
             hasError = true
-            stateLayout.infoImage(R.drawable.ic_android_black_64dp)
+            stateLayout.InfoLayoutBuilder()
+                    .infoImage(R.drawable.ic_android_black_64dp)
                     .infoTitle("Ooops.... :(")
                     .infoMessage("Unexpected error occurred. Please refresh the page!")
                     .infoButton("Refresh", onStateLayoutListener)
diff --git a/sample/src/main/res/layout/activity_state_layout_sample.xml b/sample/src/main/res/layout/activity_state_layout_sample.xml
index ae8e698..9c5393a 100644
--- a/sample/src/main/res/layout/activity_state_layout_sample.xml
+++ b/sample/src/main/res/layout/activity_state_layout_sample.xml
@@ -1,24 +1,14 @@
 <?xml version="1.0" encoding="utf-8"?>
-<android.support.constraint.ConstraintLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:app="http://schemas.android.com/apk/res-auto"
+<com.erkutaras.statelayout.StateLayout xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:tools="http://schemas.android.com/tools"
+    android:id="@+id/stateLayout"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     tools:context=".StateLayoutSampleActivity">
 
-    <com.erkutaras.statelayout.StateLayout
-        android:id="@+id/stateLayout"
+    <WebView
+        android:id="@+id/webView"
         android:layout_width="match_parent"
-        android:layout_height="match_parent"
-        android:layout_marginBottom="8dp"
-        android:layout_marginTop="8dp">
+        android:layout_height="match_parent" />
 
-        <WebView
-            android:id="@+id/webView"
-            android:layout_width="match_parent"
-            android:layout_height="match_parent"/>
-
-    </com.erkutaras.statelayout.StateLayout>
-
-</android.support.constraint.ConstraintLayout>
\ No newline at end of file
+</com.erkutaras.statelayout.StateLayout>
\ No newline at end of file