pass extra width height

app-img2img.py

@@ -21,10 +21,11 @@ import os
 import time
 import psutil
 
-
 MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
 TIMEOUT = float(os.environ.get("TIMEOUT", 0))
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
+WIDTH = 512
+HEIGHT = 512
 
 # check if MPS is available OSX only M1/M2/M3 chips
 mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
@@ -56,7 +57,7 @@ else:
         custom_revision="main",
     )
 pipe.vae = AutoencoderTiny.from_pretrained(
-    "madebyollin/taesd", torch_dtype=torch.float16, use_safetensors=True
+    "madebyollin/taesd", torch_dtype=torch_dtype, use_safetensors=True
 )
 pipe.set_progress_bar_config(disable=True)
 pipe.to(torch_device=torch_device, torch_dtype=torch_dtype).to(device)
@@ -77,18 +78,29 @@ compel_proc = Compel(
 user_queue_map = {}
 
 
-def predict(input_image, prompt, guidance_scale=8.0, strength=0.5, seed=2159232):
-    generator = torch.manual_seed(seed)
-    prompt_embeds = compel_proc(prompt)
+class InputParams(BaseModel):
+    prompt: str
+    seed: int = 2159232
+    guidance_scale: float = 8.0
+    strength: float = 0.5
+    width: int = WIDTH
+    height: int = HEIGHT
+
+
+def predict(input_image: Image.Image, params: InputParams):
+    generator = torch.manual_seed(params.seed)
+    prompt_embeds = compel_proc(params.prompt)
     # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps.
     num_inference_steps = 3
     results = pipe(
         prompt_embeds=prompt_embeds,
         generator=generator,
         image=input_image,
-        strength=strength,
+        strength=params.strength,
         num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
+        guidance_scale=params.guidance_scale,
+        width=params.width,
+        height=params.height,
         lcm_origin_steps=50,
         output_type="pil",
     )
@@ -112,13 +124,6 @@ app.add_middleware(
 )
 
 
-class InputParams(BaseModel):
-    seed: int
-    prompt: str
-    strength: float
-    guidance_scale: float
-
-
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
@@ -177,10 +182,7 @@ async def stream(user_id: uuid.UUID):
 
                 image = predict(
                     input_image,
-                    params.prompt,
-                    params.guidance_scale,
-                    params.strength,
-                    params.seed,
+                    params,
                 )
                 if image is None:
                     continue

app-txt2img.py

@@ -25,7 +25,8 @@ import psutil
 MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
 TIMEOUT = float(os.environ.get("TIMEOUT", 0))
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
-
+WIDTH = 512
+HEIGHT = 512
 # check if MPS is available OSX only M1/M2/M3 chips
 mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -66,9 +67,9 @@ pipe.unet.to(memory_format=torch.channels_last)
 if psutil.virtual_memory().total < 64 * 1024**3:
     pipe.enable_attention_slicing()
 
-if not mps_available:
-    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-    pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
+# if not mps_available:
+#     pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+#     pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
 
 compel_proc = Compel(
     tokenizer=pipe.tokenizer,
@@ -77,17 +78,25 @@ compel_proc = Compel(
 )
 user_queue_map = {}
 
-
-def predict(prompt, guidance_scale=8.0, seed=2159232):
-    generator = torch.manual_seed(seed)
-    prompt_embeds = compel_proc(prompt)
+class InputParams(BaseModel):
+    prompt: str
+    seed: int = 2159232
+    guidance_scale: float = 8.0
+    width: int = WIDTH
+    height: int = HEIGHT
+
+def predict(params: InputParams):
+    generator = torch.manual_seed(params.seed)
+    prompt_embeds = compel_proc(params.prompt)
     # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps.
     num_inference_steps = 8
     results = pipe(
         prompt_embeds=prompt_embeds,
         generator=generator,
         num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
+        guidance_scale=params.guidance_scale,
+        width=params.width,
+        height=params.height,
         lcm_origin_steps=50,
         output_type="pil",
     )
@@ -110,13 +119,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
-class InputParams(BaseModel):
-    prompt: str
-    seed: int
-    guidance_scale: float
-
-
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
@@ -173,7 +175,7 @@ async def stream(user_id: uuid.UUID):
                 if params is None:
                     continue
 
-                image = predict(params.prompt, params.guidance_scale, params.seed)
+                image = predict(params)
                 if image is None:
                     continue
                 frame_data = io.BytesIO()

img2img/index.html

@@ -10,8 +10,9 @@
     <script src="https://cdn.jsdelivr.net/npm/[email protected]/piexif.min.js"></script>
     <script src="https://cdn.tailwindcss.com"></script>
     <script type="module">
-        const WIDTH = 768;
-        const HEIGHT = 768;
+        // you can change the size of the input image to 768x768 if you have a powerful GPU
+        const WIDTH = 512;
+        const HEIGHT = 512;
         const seedEl = document.querySelector("#seed");
         const promptEl = document.querySelector("#prompt");
         const guidanceEl = document.querySelector("#guidance-scale");

requirements.txt

@@ -1,6 +1,7 @@
 diffusers==0.21.4 
 transformers==4.34.1
 gradio==3.50.2
+--extra-index-url https://download.pytorch.org/whl/cu121
 torch==2.1.0
 fastapi==0.104.0
 uvicorn==0.23.2