(
device: str = "cuda" if torch.cuda.is_available() else "cpu",
offload: bool = False,
output_dir: str = "output",
track_usage: bool = False,
)
| 56 | |
| 57 | @torch.inference_mode() |
| 58 | def main( |
| 59 | device: str = "cuda" if torch.cuda.is_available() else "cpu", |
| 60 | offload: bool = False, |
| 61 | output_dir: str = "output", |
| 62 | track_usage: bool = False, |
| 63 | ): |
| 64 | torch_device = torch.device(device) |
| 65 | names = list(configs.keys()) |
| 66 | name = st.selectbox("Which model to load?", names) |
| 67 | if name is None or not st.checkbox("Load model", False): |
| 68 | return |
| 69 | |
| 70 | is_schnell = name == "flux-schnell" |
| 71 | model, ae, t5, clip, nsfw_classifier = get_models( |
| 72 | name, |
| 73 | device=torch_device, |
| 74 | offload=offload, |
| 75 | is_schnell=is_schnell, |
| 76 | ) |
| 77 | |
| 78 | do_img2img = ( |
| 79 | st.checkbox( |
| 80 | "Image to Image", |
| 81 | False, |
| 82 | disabled=is_schnell, |
| 83 | help="Partially noise an image and denoise again to get variations.\n\nOnly works for flux-dev", |
| 84 | ) |
| 85 | and not is_schnell |
| 86 | ) |
| 87 | if do_img2img: |
| 88 | init_image = get_image() |
| 89 | if init_image is None: |
| 90 | st.warning("Please add an image to do image to image") |
| 91 | image2image_strength = st.number_input("Noising strength", min_value=0.0, max_value=1.0, value=0.8) |
| 92 | if init_image is not None: |
| 93 | h, w = init_image.shape[-2:] |
| 94 | st.write(f"Got image of size {w}x{h} ({h * w / 1e6:.2f}MP)") |
| 95 | resize_img = st.checkbox("Resize image", False) or init_image is None |
| 96 | else: |
| 97 | init_image = None |
| 98 | resize_img = True |
| 99 | image2image_strength = 0.0 |
| 100 | |
| 101 | # allow for packing and conversion to latent space |
| 102 | width = int( |
| 103 | 16 * (st.number_input("Width", min_value=128, value=1360, step=16, disabled=not resize_img) // 16) |
| 104 | ) |
| 105 | height = int( |
| 106 | 16 * (st.number_input("Height", min_value=128, value=768, step=16, disabled=not resize_img) // 16) |
| 107 | ) |
| 108 | num_steps = int(st.number_input("Number of steps", min_value=1, value=(4 if is_schnell else 50))) |
| 109 | guidance = float(st.number_input("Guidance", min_value=1.0, value=3.5, disabled=is_schnell)) |
| 110 | seed_str = st.text_input("Seed", disabled=is_schnell) |
| 111 | if seed_str.isdecimal(): |
| 112 | seed = int(seed_str) |
| 113 | else: |
| 114 | st.info("No seed set, set to positive integer to enable") |
| 115 | seed = None |
nothing calls this directly
no test coverage detected
searching dependent graphs…