import gradio as gr import config from inference import DiffusionInference from controlnet_pipeline import ControlNetPipeline from PIL import Image import io # Initialize the inference class inference = DiffusionInference() # Initialize the ControlNet pipeline controlnet = ControlNetPipeline() def text_to_image_fn(prompt, model, negative_prompt=None, guidance_scale=7.5, num_inference_steps=50, seed=None): try: # Model validation - fallback to default if empty if not model or model.strip() == '': model = config.DEFAULT_TEXT2IMG_MODEL # Prepare seed parameter seed_value = None if seed and seed.strip() != '': try: seed_value = int(seed) except (ValueError, TypeError): # Let inference handle invalid seed pass # Create kwargs dictionary for parameters kwargs = { "prompt": prompt, "model_name": model, "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps, "seed": seed_value } # Only add negative_prompt if it's not None if negative_prompt is not None: kwargs["negative_prompt"] = negative_prompt # Call the inference module with unpacked kwargs image = inference.text_to_image(**kwargs) if image is None: return None, "No image was generated. Check the model and parameters." return image, None except Exception as e: error_msg = f"Error: {str(e)}" print(error_msg) return None, error_msg def image_to_image_fn(image, prompt, model, use_controlnet=False, negative_prompt=None, guidance_scale=7.5, num_inference_steps=50): """ Handle image to image transformation request """ if image is None: return None, "No input image provided." # Handle empty prompt - use default if completely empty if prompt is None or prompt.strip() == "": prompt = config.DEFAULT_IMG2IMG_PROMPT try: if use_controlnet: # Use ControlNet pipeline directly on the device result = controlnet.generate( prompt=prompt, image=image, negative_prompt=negative_prompt, guidance_scale=float(guidance_scale), num_inference_steps=int(num_inference_steps) ) return result, None else: # Model validation - fallback to default if empty if not model or model.strip() == '': model = config.DEFAULT_IMG2IMG_MODEL # Use regular inference API result = inference.image_to_image( image=image, prompt=prompt, model_name=model, negative_prompt=negative_prompt, guidance_scale=float(guidance_scale) if guidance_scale is not None else None, num_inference_steps=int(num_inference_steps) if num_inference_steps is not None else None ) if result is None: return None, "No image was generated. Check the model and parameters." return result, None except Exception as e: error_msg = f"Error: {str(e)}" print(error_msg) print(f"Input image type: {type(image)}") print(f"Prompt: {prompt}") print(f"Model: {model}") return None, error_msg # Create Gradio UI with gr.Blocks(title="Diffusion Models") as app: gr.Markdown("# Hugging Face Diffusion Models") with gr.Tab("Text to Image"): with gr.Row(): with gr.Column(): txt2img_prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", value=config.DEFAULT_TEXT2IMG_PROMPT) txt2img_negative = gr.Textbox(label="Negative Prompt (Optional)", placeholder="What to exclude from the image", value=config.DEFAULT_NEGATIVE_PROMPT) txt2img_model = gr.Textbox(label="Model", placeholder=f"Enter model name", value=config.DEFAULT_TEXT2IMG_MODEL) txt2img_guidance = gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.5, label="Guidance Scale") txt2img_steps = gr.Slider(minimum=10, maximum=100, value=50, step=1, label="Inference Steps") txt2img_seed = gr.Textbox(label="Seed (Optional)", placeholder="Leave empty for random seed", value="") txt2img_button = gr.Button("Generate Image") with gr.Column(): txt2img_output = gr.Image(type="pil", label="Generated Image") txt2img_error = gr.Textbox(label="Error", visible=True) txt2img_button.click( fn=text_to_image_fn, inputs=[txt2img_prompt, txt2img_model, txt2img_negative, txt2img_guidance, txt2img_steps, txt2img_seed], outputs=[txt2img_output, txt2img_error] ) with gr.Tab("Image to Image"): with gr.Row(): with gr.Column(): img2img_input = gr.Image(type="pil", label="Input Image") img2img_prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", value=config.DEFAULT_IMG2IMG_PROMPT) img2img_negative = gr.Textbox(label="Negative Prompt (Optional)", placeholder="What to exclude from the image", value=config.DEFAULT_NEGATIVE_PROMPT) with gr.Row(): with gr.Column(scale=1): img2img_controlnet = gr.Checkbox(label="Use ControlNet (Depth)", value=False) with gr.Column(scale=2): img2img_model = gr.Textbox(label="Model (used only if ControlNet is disabled)", placeholder=f"Enter model name", value=config.DEFAULT_IMG2IMG_MODEL, visible=True) img2img_guidance = gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.5, label="Guidance Scale") img2img_steps = gr.Slider(minimum=10, maximum=100, value=50, step=1, label="Inference Steps") img2img_button = gr.Button("Transform Image") with gr.Column(): img2img_output = gr.Image(type="pil", label="Generated Image") img2img_error = gr.Textbox(label="Error", visible=True) img2img_button.click( fn=image_to_image_fn, inputs=[img2img_input, img2img_prompt, img2img_model, img2img_controlnet, img2img_negative, img2img_guidance, img2img_steps], outputs=[img2img_output, img2img_error] ) # Add visibility toggle for the model textbox based on ControlNet checkbox def toggle_model_visibility(use_controlnet): return not use_controlnet img2img_controlnet.change( fn=toggle_model_visibility, inputs=[img2img_controlnet], outputs=[img2img_model] ) # Launch the Gradio app if __name__ == "__main__": app.launch(server_name=config.GRADIO_HOST, server_port=config.GRADIO_PORT)