()
| 47 | |
| 48 | |
| 49 | def main(): |
| 50 | args = get_args() |
| 51 | # Note: This will load the models, which can take several minutes |
| 52 | sd = StableDiffusionMGX(args.pipeline_type, args.onnx_model_path, |
| 53 | args.compiled_model_path, args.use_refiner, |
| 54 | args.refiner_onnx_model_path, |
| 55 | args.refiner_compiled_model_path, args.fp16, |
| 56 | args.force_compile, args.exhaustive_tune) |
| 57 | sd.warmup(5) |
| 58 | |
| 59 | def gr_wrapper(prompt, negative_prompt, steps, seed, scale, refiner_steps, |
| 60 | aesthetic_score, negative_aesthetic_score): |
| 61 | img = None |
| 62 | try: |
| 63 | oldStdout, oldStderr = sys.stdout, sys.stderr |
| 64 | sys.stdout, sys.stderr = PrintWrapper(sys.stdout), PrintWrapper( |
| 65 | sys.stderr) |
| 66 | result = sd.run( |
| 67 | str(prompt), |
| 68 | str(negative_prompt), |
| 69 | int(steps), |
| 70 | int(seed), |
| 71 | float(scale), |
| 72 | int(refiner_steps), |
| 73 | float(aesthetic_score), |
| 74 | float(negative_aesthetic_score), |
| 75 | args.verbose, |
| 76 | ) |
| 77 | img = StableDiffusionMGX.convert_to_rgb_image(result) |
| 78 | sd.print_summary(steps) |
| 79 | finally: |
| 80 | log = ''.join([sys.stdout.get_log(), sys.stderr.get_log()]) |
| 81 | sys.stdout, sys.stderr = oldStdout, oldStderr |
| 82 | return img, log |
| 83 | |
| 84 | demo = gr.Interface(gr_wrapper, [ |
| 85 | gr.Textbox(value=args.prompt, label="Prompt"), |
| 86 | gr.Textbox(value=args.negative_prompt, |
| 87 | label="Negative prompt (Optional)"), |
| 88 | gr.Slider(1, 100, step=1, value=args.steps, label="Number of steps"), |
| 89 | gr.Textbox(value=args.seed, label="Random seed"), |
| 90 | gr.Slider(1, 20, step=0.1, value=args.scale, label="Guidance scale"), |
| 91 | gr.Slider(0, |
| 92 | 100, |
| 93 | step=1, |
| 94 | value=args.refiner_steps, |
| 95 | label="Number of refiner steps. (Use 0 to skip it)", |
| 96 | visible=args.use_refiner), |
| 97 | gr.Slider(1, |
| 98 | 20, |
| 99 | step=0.1, |
| 100 | value=args.refiner_aesthetic_score, |
| 101 | label="Aesthetic score (Refiner)", |
| 102 | visible=args.use_refiner), |
| 103 | gr.Slider(1, |
| 104 | 20, |
| 105 | step=0.1, |
| 106 | value=args.refiner_negative_aesthetic_score, |
no test coverage detected