update requirements

app.py

@@ -9,15 +9,30 @@ import spaces
 import torch
 import random
 from PIL import Image
-
+import torchvision.transforms.functional as TF
 from pipeline import GenSIRR
 from diffusers.utils import load_image
-
+import torch.nn.functional as F
 from optimization import optimize_pipeline_
 
 MAX_SEED = np.iinfo(np.int32).max
 from huggingface_hub import hf_hub_download
 
+def pad_for_model(image: torch.Tensor, multiple: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
+    """Pad the tensor image so height/width are divisible by ``multiple``."""
+    if multiple <= 0:
+        raise ValueError("round_multiple must be > 0")
+
+    height, width = image.shape[-2:]
+    pad_h = (multiple - height % multiple) % multiple
+    pad_w = (multiple - width % multiple) % multiple
+    if pad_h == 0 and pad_w == 0:
+        return image, (0, 0)
+
+    padded = F.pad(image.unsqueeze(0), (0, pad_w, 0, pad_h), mode="reflect").squeeze(0)
+    return padded, (pad_h, pad_w)
+
+
 def load_deepspeed_weights(model, checkpoint_path) -> None:
     """Load LoRA weights from a DeepSpeed ZeRO Stage 2 checkpoint into the model."""
     tensor_path = checkpoint_path
@@ -81,10 +96,20 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-
+    size = 768
     input_image = input_image.convert("RGB")
+    if input_image.width < input_image.height:
+        input_image = input_image.resize((size, int(size * input_image.height / input_image.width)))
+    else:
+        input_image = input_image.resize((int(size * input_image.width / input_image.height), size))
+    tensor = TF.to_tensor(img)
+
+    original_size = tensor.shape[-2:]
+    padded_tensor, padding = pad_for_model(tensor, args.round_multiple)
+    batch_cpu = padded_tensor.unsqueeze(0)
+    batch_device = batch_cpu.to('cuda')
     image = pipe(
-        image=input_image, 
+        image=batch_device, 
         width = input_image.size[0],
         height = input_image.size[1],
         num_inference_steps=steps,

requirements.txt

@@ -3,4 +3,5 @@ git+https://github.com/huggingface/diffusers.git
 accelerate
 safetensors
 sentencepiece
-peft
+peft
+torchvision