Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| import numpy as np | |
| import torch, random, json, spaces, time | |
| # from safetensors.torch import load_file | |
| # from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler | |
| # from videox_fun.pipeline import ZImageControlPipeline | |
| # from videox_fun.models import ZImageControlTransformer2DModel | |
| # from transformers import AutoTokenizer, Qwen3ForCausalLM | |
| # from diffusers import AutoencoderKL | |
| # from controlnet_aux.processor import Processor | |
| from utils import repo_utils, image_utils, prompt_utils | |
| repo_utils.clone_repo_if_not_exists("https://github.com/apple/ml-starflow.git", "app/models") | |
| repo_utils.clone_repo_if_not_exists("https://huggingface.co/apple/starflow", "app/models") | |
| # MODEL_PATH = "models/Z-Image-Turbo/" | |
| # CONTROLNET_PATH = "models/Z-Image-Turbo-Fun-Controlnet-Union/Z-Image-Turbo-Fun-Controlnet-Union.safetensors" | |
| DTYPE = torch.bfloat16 | |
| MAX_SEED = np.iinfo(np.int32).max | |
| # # load transformer | |
| # transformer = ZImageControlTransformer2DModel.from_pretrained( | |
| # MODEL_PATH, | |
| # subfolder="transformer", | |
| # transformer_additional_kwargs={ | |
| # "control_layers_places": [0, 5, 10, 15, 20, 25], | |
| # "control_in_dim": 16 | |
| # }, | |
| # torch_dtype= DTYPE | |
| # ).to("cuda") | |
| # ## Load controlnet | |
| # state_dict = load_file(CONTROLNET_PATH) | |
| # state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict | |
| # m, u = transformer.load_state_dict(state_dict, strict=False) | |
| # print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") | |
| # # load ZImageControlPipeline | |
| # vae = AutoencoderKL.from_pretrained( | |
| # MODEL_PATH, | |
| # subfolder="vae", | |
| # device_map="cuda", | |
| # torch_dtype= DTYPE | |
| # ) | |
| # tokenizer = AutoTokenizer.from_pretrained( | |
| # MODEL_PATH, | |
| # subfolder="tokenizer" | |
| # ) | |
| # text_encoder = Qwen3ForCausalLM.from_pretrained( | |
| # MODEL_PATH, | |
| # subfolder="text_encoder", | |
| # torch_dtype=DTYPE, | |
| # ) | |
| # scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( | |
| # MODEL_PATH, | |
| # subfolder="scheduler" | |
| # ) | |
| # pipe = ZImageControlPipeline( | |
| # vae=vae, | |
| # tokenizer=tokenizer, | |
| # text_encoder=text_encoder, | |
| # transformer=transformer, | |
| # scheduler=scheduler, | |
| # ) | |
| # pipe.to("cuda", DTYPE) | |
| # def prepare(prompt, is_polish_prompt): | |
| # if not is_polish_prompt: return prompt, False | |
| # polished_prompt = prompt_utils.polish_prompt(prompt) | |
| # return polished_prompt, True | |
| def inference( | |
| prompt, | |
| negative_prompt, | |
| seed=42, | |
| randomize_seed=True, | |
| guidance_scale=1.5, | |
| num_inference_steps=8, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| timestamp = time.time() | |
| print(f"timestamp: {timestamp}") | |
| # # process image | |
| # print("DEBUG: process image") | |
| # if input_image is None: | |
| # print("Error: input_image is empty.") | |
| # return None | |
| # print("DEBUG: control_image_torch") | |
| # orig_width, orig_height = input_image.size | |
| # control_image, width, height = image_utils.rescale_image(input_image, image_scale, 16, 2048) | |
| # control_image_torch = image_utils.get_image_latent(control_image, sample_size=[height, width])[:, :, 0] | |
| # # generation | |
| # if randomize_seed: seed = random.randint(0, MAX_SEED) | |
| # generator = torch.Generator().manual_seed(seed) | |
| # output_image = pipe( | |
| # prompt=prompt, | |
| # negative_prompt = negative_prompt, | |
| # width=width, | |
| # height=height, | |
| # generator=generator, | |
| # guidance_scale=guidance_scale, | |
| # control_image=control_image_torch, | |
| # num_inference_steps=num_inference_steps, | |
| # control_context_scale=control_context_scale, | |
| # ).images[0] | |
| # output_image = output_image.resize((orig_width * image_scale, orig_height * image_scale)) | |
| # return output_image, seed | |
| def read_file(path: str) -> str: | |
| with open(path, 'r', encoding='utf-8') as f: | |
| content = f.read() | |
| return content | |
| css = """ | |
| #col-container { | |
| margin: 0 auto; | |
| max-width: 960px; | |
| } | |
| """ | |
| with open('examples/0_examples.json', 'r') as file: examples = json.load(file) | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Column(): | |
| gr.HTML(read_file("static/header.html")) | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| show_label=False, | |
| lines=2, | |
| placeholder="Enter your prompt", | |
| value="a man in a fishing boat. high quality, detailed" | |
| # container=False, | |
| ) | |
| # is_polish_prompt = gr.Checkbox(label="Polish prompt", value=True) | |
| # control_mode = gr.Radio( | |
| # choices=["Canny", "Depth", "HED", "MLSD", "Pose"], | |
| # value="Canny", | |
| # label="Control Mode" | |
| # ) | |
| run_button = gr.Button("Generate", variant="primary") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| negative_prompt = gr.Textbox( | |
| label="Negative prompt", | |
| lines=2, | |
| container=False, | |
| placeholder="Enter your negative prompt", | |
| value="blurry, ugly, bad" | |
| ) | |
| with gr.Row(): | |
| num_inference_steps = gr.Slider( | |
| label="Steps", | |
| minimum=1, | |
| maximum=30, | |
| step=1, | |
| value=9, | |
| ) | |
| control_context_scale = gr.Slider( | |
| label="Context scale", | |
| minimum=0.0, | |
| maximum=1.0, | |
| step=0.01, | |
| value=0.75, | |
| ) | |
| with gr.Row(): | |
| guidance_scale = gr.Slider( | |
| label="Guidance scale", | |
| minimum=0.0, | |
| maximum=10.0, | |
| step=0.1, | |
| value=1.0, | |
| ) | |
| seed = gr.Slider( | |
| label="Seed", | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=42, | |
| ) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=False) | |
| with gr.Column(): | |
| output_image = gr.Image(label="Generated image", show_label=False) | |
| # polished_prompt = gr.Textbox(label="Polished prompt", interactive=False) | |
| # with gr.Accordion("Preprocessor output", open=False): | |
| # control_image = gr.Image(label="Control image", show_label=False) | |
| # gr.Examples(examples=examples, inputs=[input_image]) | |
| gr.Markdown(read_file("static/footer.md")) | |
| run_button.click( | |
| fn=inference, | |
| inputs=[ | |
| prompt, | |
| negative_prompt, | |
| seed, | |
| randomize_seed, | |
| guidance_scale, | |
| num_inference_steps, | |
| ], | |
| outputs=[output_image, seed], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True, css=css) | |