baxin/isometric-map-generator
0
1import gradio as gr2import numpy as np3import random4# import spaces #[uncomment to use ZeroGPU]5from diffusers import DiffusionPipeline6import torch7 8# --- Model and Device Configuration ---9 10# Global dictionary to cache loaded models, preventing re-loading.11pipelines = {}12# Mapping of user-friendly names to Hugging Face model repository IDs.13MODEL_MAP = {14 "SDXL-Turbo": "stabilityai/sdxl-turbo",15 "Nano-Banana": "emilianJR/nano-banana-base-1.0"16}17 18device = "cuda" if torch.cuda.is_available() else "cpu"19torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float3220 21# This function loads a model if it's not already in our cache22def get_pipeline(model_name: str):23 """Loads and caches a diffusion pipeline based on the model name."""24 repo_id = MODEL_MAP[model_name]25 if repo_id not in pipelines:26 print(f"Loading model: {repo_id}...")27 pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch_dtype, variant="fp16" if torch.cuda.is_available() else "fp32")28 pipe.to(device)29 pipelines[repo_id] = pipe30 print("Model loaded successfully.")31 return pipelines[repo_id]32 33MAX_SEED = np.iinfo(np.int32).max34MAX_IMAGE_SIZE = 102435 36# --- Inference Function ---37 38# @spaces.GPU #[uncomment to use ZeroGPU]39def infer(40 prompt,41 negative_prompt,42 model_selection, # New parameter to select the model43 seed,44 randomize_seed,45 width,46 height,47 guidance_scale,48 num_inference_steps,49 progress=gr.Progress(track_tqdm=True),50):51 # Load the selected pipeline52 pipe = get_pipeline(model_selection)53 54 if randomize_seed:55 seed = random.randint(0, MAX_SEED)56 57 generator = torch.Generator(device=device).manual_seed(seed)58 59 # SDXL-Turbo does not use guidance_scale, so we set it to 0.0 if that model is selected.60 # Other models might need it.61 effective_guidance_scale = 0.0 if model_selection == "SDXL-Turbo" else guidance_scale62 63 image = pipe(64 prompt=prompt,65 negative_prompt=negative_prompt,66 guidance_scale=effective_guidance_scale,67 num_inference_steps=num_inference_steps,68 width=width,69 height=height,70 generator=generator,71 ).images[0]72 73 return image, seed74 75# --- UI Helper Function ---76 77def update_settings_for_model(model_selection: str):78 """Updates the UI with recommended settings for the chosen model."""79 if model_selection == "SDXL-Turbo":80 # SDXL-Turbo works best with low steps and no guidance81 return gr.Slider(value=0.0), gr.Slider(value=2)82 elif model_selection == "Nano-Banana":83 # A more standard SDXL setup84 return gr.Slider(value=7.5), gr.Slider(value=25)85 return gr.Slider(), gr.Slider() # Default empty update86 87# --- Gradio UI Layout ---88 89examples = [90 "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",91 "An astronaut riding a green horse",92 "A delicious ceviche cheesecake slice",93]94 95css = """96#col-container {97 margin: 0 auto;98 max-width: 640px;99}100"""101 102with gr.Blocks(css=css) as demo:103 with gr.Column(elem_id="col-container"):104 gr.Markdown("# Text-to-Image with Model Switching")105 106 with gr.Row():107 prompt = gr.Text(108 label="Prompt",109 show_label=False,110 max_lines=1,111 placeholder="Enter your prompt",112 container=False,113 )114 run_button = gr.Button("Run", scale=0, variant="primary")115 116 model_selection = gr.Radio(117 label="Select Model",118 choices=list(MODEL_MAP.keys()),119 value="SDXL-Turbo",120 )121 122 result = gr.Image(label="Result", show_label=False, type="pil")123 124 with gr.Accordion("Advanced Settings", open=False):125 # 1. Added Gemini API Key input box126 gemini_api_key = gr.Textbox(127 label="Gemini API Key",128 placeholder="Enter your Gemini API key here",129 type="password",130 visible=True, # Set to True to make it visible131 )132 negative_prompt = gr.Text(133 label="Negative prompt",134 max_lines=1,135 placeholder="Enter a negative prompt",136 )137 seed = gr.Slider(138 label="Seed",139 minimum=0,140 maximum=MAX_SEED,141 step=1,142 value=0,143 )144 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)145 with gr.Row():146 width = gr.Slider(147 label="Width",148 minimum=256,149 maximum=MAX_IMAGE_SIZE,150 step=32,151 value=512, # Changed default to 512 for SDXL-Turbo152 )153 height = gr.Slider(154 label="Height",155 minimum=256,156 maximum=MAX_IMAGE_SIZE,157 step=32,158 value=512, # Changed default to 512 for SDXL-Turbo159 )160 with gr.Row():161 guidance_scale = gr.Slider(162 label="Guidance scale",163 minimum=0.0,164 maximum=20.0,165 step=0.1,166 value=0.0, # Default for SDXL-Turbo167 )168 num_inference_steps = gr.Slider(169 label="Number of inference steps",170 minimum=1,171 maximum=50,172 step=1,173 value=2, # Default for SDXL-Turbo174 )175 gr.Examples(examples=examples, inputs=[prompt])176 177 # --- Event Handlers ---178 179 # Main inference trigger180 gr.on(181 triggers=[run_button.click, prompt.submit],182 fn=infer,183 inputs=[184 prompt,185 negative_prompt,186 model_selection,187 seed,188 randomize_seed,189 width,190 height,191 guidance_scale,192 num_inference_steps,193 ],194 outputs=[result, seed],195 )196 197 # Trigger to update settings when the model selection changes198 model_selection.change(199 fn=update_settings_for_model,200 inputs=model_selection,201 outputs=[guidance_scale, num_inference_steps]202 )203 204 205if __name__ == "__main__":206 demo.launch(debug=True)