app.py

import gradio as gr
import config
from inference import DiffusionInference
from PIL import Image
import io

# Initialize the inference class
inference = DiffusionInference()

def text_to_image_fn(prompt, model, negative_prompt=None, guidance_scale=7.5, num_inference_steps=50):
    """
    Handle text to image generation request
    """
    try:
        if not model:
            model = config.DEFAULT_TEXT2IMG_MODEL
            
        # Call the inference module
        image = inference.text_to_image(
            prompt=prompt,
            model_name=model,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps
        )
        
        return image, None
    except Exception as e:
        return None, str(e)

def image_to_image_fn(image, prompt, model, negative_prompt=None, guidance_scale=7.5, num_inference_steps=50):
    """
    Handle image to image transformation request
    """
    try:
        if not model:
            model = config.DEFAULT_IMG2IMG_MODEL
            
        # Call the inference module
        result = inference.image_to_image(
            image=image,
            prompt=prompt,
            model_name=model,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps
        )
        
        return result, None
    except Exception as e:
        return None, str(e)

# 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...")
                txt2img_negative = gr.Textbox(label="Negative Prompt (Optional)", placeholder="What to exclude from the image")
                txt2img_model = gr.Textbox(label="Model", placeholder=f"Enter model name (default: {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_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],
            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...")
                img2img_negative = gr.Textbox(label="Negative Prompt (Optional)", placeholder="What to exclude from the image")
                img2img_model = gr.Textbox(label="Model", placeholder=f"Enter model name (default: {config.DEFAULT_IMG2IMG_MODEL})")
                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_negative, img2img_guidance, img2img_steps],
            outputs=[img2img_output, img2img_error]
        )

# Launch the Gradio app
if __name__ == "__main__":
    app.launch(server_name=config.GRADIO_HOST, server_port=config.GRADIO_PORT)