From fe4632034b1ffcd46e9347b03b76e8c51104b14a Mon Sep 17 00:00:00 2001
From: Stephan Auerhahn <stephan@stability.ai>
Date: Fri, 11 Aug 2023 16:31:53 -0700
Subject: [PATCH] fix for orig dimensions

---
 scripts/demo/streamlit_helpers.py | 39 +++++++++++++++----------------
 sgm/inference/api.py              |  6 +++--
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/scripts/demo/streamlit_helpers.py b/scripts/demo/streamlit_helpers.py
index 84c4e62..93293dd 100644
--- a/scripts/demo/streamlit_helpers.py
+++ b/scripts/demo/streamlit_helpers.py
@@ -32,28 +32,27 @@ def init_st(
     enable_swap=True,
 ) -> Dict[str, Any]:
     state: Dict[str, Any] = dict()
-    if not "model" in state:
-        config = spec.config
-        ckpt = spec.ckpt
+    config = spec.config
+    ckpt = spec.ckpt
 
-        if enable_swap:
-            pipeline = SamplingPipeline(
-                model_spec=spec,
-                use_fp16=use_fp16,
-                device=CudaModelManager(device="cuda", swap_device="cpu"),
-            )
-        else:
-            pipeline = SamplingPipeline(model_spec=spec, use_fp16=use_fp16)
+    if enable_swap:
+        pipeline = SamplingPipeline(
+            model_spec=spec,
+            use_fp16=use_fp16,
+            device=CudaModelManager(device="cuda", swap_device="cpu"),
+        )
+    else:
+        pipeline = SamplingPipeline(model_spec=spec, use_fp16=use_fp16)
 
-        state["spec"] = spec
-        state["model"] = pipeline
-        state["ckpt"] = ckpt if load_ckpt else None
-        state["config"] = config
-        state["params"] = SamplingParams()
-        if load_filter:
-            state["filter"] = DeepFloydDataFiltering(verbose=False)
-        else:
-            state["filter"] = None
+    state["spec"] = spec
+    state["model"] = pipeline
+    state["ckpt"] = ckpt if load_ckpt else None
+    state["config"] = config
+    state["params"] = SamplingParams()
+    if load_filter:
+        state["filter"] = DeepFloydDataFiltering(verbose=False)
+    else:
+        state["filter"] = None
     return state
 
 
diff --git a/sgm/inference/api.py b/sgm/inference/api.py
index e680dc5..eccf129 100644
--- a/sgm/inference/api.py
+++ b/sgm/inference/api.py
@@ -75,8 +75,8 @@ class SamplingParams:
     aesthetic_score: float = 6.0
     negative_aesthetic_score: float = 2.5
     img2img_strength: float = 1.0
-    orig_width: int = 1024
-    orig_height: int = 1024
+    orig_width: int = width
+    orig_height: int = height
     crop_coords_top: int = 0
     crop_coords_left: int = 0
     sigma_min: float = 0.0292
@@ -301,6 +301,8 @@ class SamplingPipeline:
         value_dict["negative_prompt"] = negative_prompt
         value_dict["target_width"] = width
         value_dict["target_height"] = height
+        value_dict["orig_width"] = width
+        value_dict["orig_height"] = height
         return do_img2img(
             image,
             self.model,