ovi054/FLUX-Dev-Replicate-API

1import os2import io3import random4import requests5import gradio as gr6import numpy as np7from PIL import Image8import replicate9 10 11MAX_SEED = np.iinfo(np.int32).max12 13 14def predict(replicate_api, prompt, lora_id, lora_scale=0.95, aspect_ratio="1:1", seed=-1, randomize_seed=True, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):15 16    # Validate API key and prompt17    if not replicate_api or not prompt:18        return "Error: Missing necessary inputs.", -1, None19    20    # Set the seed if randomize_seed is True21    if randomize_seed:22        seed = random.randint(0, MAX_SEED)23 24    # Set the Replicate API token in the environment variable25    os.environ["REPLICATE_API_TOKEN"] = replicate_api26 27    # Construct the input for the replicate model28    input_params = {29        "prompt": prompt,30        "output_format": "jpg",31        "aspect_ratio": aspect_ratio,32        "num_inference_steps": num_inference_steps,33        "guidance_scale": guidance_scale,34        "seed": seed,35        "disable_safety_checker": True36    }37 38    # If lora_id is provided, include it in the input39    if lora_id and lora_id.strip()!="":40        input_params["hf_lora"] = lora_id.strip()41        input_params["lora_scale"] = lora_scale42 43    try:44        # Run the model using the user's API token from the environment variable45        output = replicate.run(46            "lucataco/flux-dev-lora:a22c463f11808638ad5e2ebd582e07a469031f48dd567366fb4c6fdab91d614d",47            input=input_params48        )49        print("\nGeneration Completed: ",output,prompt,lora_id)50        return output[0], seed, seed  # Return the generated image and seed51 52    except Exception as e:53        # Catch any exceptions, such as invalid API token or lack of credits54        return f"Error: {str(e)}", -1, None55 56    finally:57        # Always remove the API key from the environment58        if "REPLICATE_API_TOKEN" in os.environ:59            del os.environ["REPLICATE_API_TOKEN"]60 61    62 63demo = gr.Interface(fn=predict, inputs="text", outputs="image")64 65css="""66#col-container {67    margin: 0 auto;68    max-width: 520px;69}70"""71 72examples = [73    "a tiny astronaut hatching from an egg on the moon",74    "a cat holding a sign that says hello world",75    "an anime illustration of a wiener schnitzel",76]77 78with gr.Blocks(css=css) as demo:79    with gr.Column(elem_id="col-container"):80        gr.Markdown("# FLUX Dev with Replicate API")81        82        replicate_api = gr.Text(label="Replicate API Key", type='password', show_label=True, max_lines=1, placeholder="Enter your Replicate API token", container=True)83        prompt = gr.Text(label="Prompt", show_label=True, lines = 2, max_lines=4, show_copy_button = True, placeholder="Enter your prompt", container=True)84        with gr.Accordion("Advanced Settings", open=False):85            with gr.Row():86                custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path (optional)", placeholder="multimodalart/vintage-ads-flux")87                lora_scale = gr.Slider(88                    label="LoRA Scale",89                    minimum=0,90                    maximum=1,91                    step=0.01,92                    value=0.95,93                )94            aspect_ratio = gr.Radio(label="Aspect ratio", value="1:1", choices=["1:1", "4:5", "2:3", "3:4","9:16", "4:3", "16:9"])95            seed = gr.Slider(96                label="Seed",97                minimum=0,98                maximum=MAX_SEED,99                step=1,100                value=0,101            )102 103            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)104 105            with gr.Row():106                guidance_scale = gr.Slider(107                    label="Guidance Scale",108                    minimum=1,109                    maximum=15,110                    step=0.1,111                    value=3.5,112                )113                num_inference_steps = gr.Slider(114                    label="Number of inference steps",115                    minimum=1,116                    maximum=50,117                    step=1,118                    value=28,119                )120        submit = gr.Button("Generate Image", variant="primary",scale=1)121 122        output = gr.Image(label="Output Image", show_label=True)123 124        seed_used = gr.Textbox(label="Seed Used", show_copy_button = True)125        126 127        gr.Examples(128            examples=examples,129            fn=predict,130            inputs=[prompt]131        )132        gr.on(133            triggers=[submit.click, prompt.submit],134            fn=predict,135            inputs=[replicate_api, prompt, custom_lora, lora_scale, aspect_ratio, seed, randomize_seed, guidance_scale, num_inference_steps],136            outputs = [output, seed, seed_used]137        )138 139demo.launch()