add more params for replicate

predict.py

@@ -6,6 +6,7 @@ from cog import BasePredictor, Input, Path
 import os
 import sys
 import signal
+import time
 import re
 from typing import Dict, List, Any
 
@@ -75,13 +76,50 @@ def initialize():
     shared.opts.onchange("temp_dir", ui_tempdir.on_tmpdir_changed)
 
     # make the program just exit at ctrl+c without waiting for anything
-    def sigint_handler(sig, frame):
-        print(f'Interrupted with signal {sig} in {frame}')
-        os._exit(0)
+    # def sigint_handler(sig, frame):
+    #     print(f'Interrupted with signal {sig} in {frame}')
+    #     os._exit(0)
 
-    signal.signal(signal.SIGINT, sigint_handler)
+    # signal.signal(signal.SIGINT, sigint_handler)
 
 class Predictor(BasePredictor):
+    def setup(self):
+        """Load the model into memory to make running multiple predictions efficient"""
+        initialize()
+
+    def predict(
+        self,
+        prompt: str = Input(description="prompt en", default="lora:koreanDollLikeness_v15:0.66, best quality, ultra high res, (photorealistic:1.4), 1girl, beige sweater, black choker, smile, laughing, bare shoulders, solo focus, ((full body), (brown hair:1), looking at viewer"),
+        negative_prompt: str = Input(description="negative prompt", default="paintings, sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, age spot, glans, (ugly:1.331), (duplicate:1.331), (morbid:1.21), (mutilated:1.21), (tranny:1.331), mutated hands, (poorly drawn hands:1.331), blurry, 3hands,4fingers,3arms, bad anatomy, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts,poorly drawn face,mutation,deformed"),
+        sampler_name: str = Input(description="sampler name", default="DPM++ SDE Karras", choices=["DPM++ SDE Karras", "DPM++ 2M Karras", "DPM++ 2S a Karras", "DPM2 a Karras", "DPM2 Karras", "LMS Karras", "DPM adaptive", "DPM fast", "DPM++ SDE", "DPM++ 2M", "DPM++ 2S a", "DPM2 a", "DPM2", "Heun", "LMS", "Euler", "Euler a"]),
+        steps: int = Input(description="steps", default=20),
+        cfg_scale: int = Input(description="cfg scale", default=8),
+        width: int = Input(description="width", default=512),
+        height: int = Input(description="height", default=768),
+        seed: int = Input(description="seed", default=-1),
+    ) -> Path:
+        """Run a single prediction on the model"""
+        args = {
+            "do_not_save_samples": True,
+            "do_not_save_grid": True,
+            "outpath_samples": "./output",
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "sampler_name": sampler_name,
+            "steps": steps, # 25
+            "cfg_scale": cfg_scale,
+            "width": width,
+            "height": height,
+            "seed": seed,
+        }
+        p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)
+        processed = process_images(p)
+        filename = str(int(time.time())) + ".png"
+        processed.images[0].save(fp=filename, format="PNG")
+        # single_image_b64 = encode_pil_to_base64(processed.images[0]).decode('utf-8')
+        return Path(filename)
+
+class PredictorOld(BasePredictor):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
         initialize()