File size: 5,715 Bytes
aec26a9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import gradio as gr
import gradio as gr
from groq import Groq
import os
from deep_translator import GoogleTranslator
from deep_translator import GoogleTranslator  # Import the GoogleTranslator class
import whisper
import gradio as gr
from groq import Groq
import os
from deep_translator import GoogleTranslator # Import the GoogleTranslator class
import pickle
import whisper
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
import matplotlib.pyplot as plt
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file


# Replace with your actual API key
api_key = "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"
client = Groq(api_key=api_key)

# Load the custom model for image generation
base = "stabilityai/stable-diffusion-xl-base-1.0"
repo = "ByteDance/SDXL-Lightning"
ckpt = "sdxl_lightning_4step_unet.safetensors"  # Ensure the correct checkpoint

# Load the custom UNet and set up the pipeline
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cpu", torch.float16)
unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cpu"))
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cpu")
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")


# Function to transcribe, translate, and generate an image
def process_audio(audio_path, generate_image):
    if audio_path is None:
        return "Please upload an audio file.", None, None

    # Step 1: Transcribe audio
    try:
        with open(audio_path, "rb") as file:
            transcription = client.audio.transcriptions.create(
                file=(os.path.basename(audio_path), file.read()),
                model="whisper-large-v3",
                language="ta",
                response_format="verbose_json",
            )
        tamil_text = transcription.text
    except Exception as e:
        return f"An error occurred during transcription: {str(e)}", None, None

    # Step 2: Translate Tamil to English
    try:
        translator = GoogleTranslator(source='ta', target='en')
        translation = translator.translate(tamil_text)
    except Exception as e:
        return tamil_text, f"An error occurred during translation: {str(e)}", None

    # Step 3: Generate image (if selected)
    if generate_image:
        try:
            # Use the custom model and pipeline to generate an image
            img = pipe(translation, num_inference_steps=4, guidance_scale=0).images[0]
            return tamil_text, translation, img
        except Exception as e:
            return tamil_text, translation, f"An error occurred during image generation: {str(e)}"

    return tamil_text, translation, None


# Function for direct prompt to image generation
def generate_image_from_prompt(prompt):
    try:
        img = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
        return img
    except Exception as e:
        return f"An error occurred during image generation: {str(e)}"


# Assuming your 'process_audio' and 'generate_image_from_prompt' functions are defined elsewhere

# Gradio interface with the requested customizations
with gr.Blocks(css="""

    .gradio-container {background-color: #D8D2C2;} 

    .btn-red {background-color: red; color: white;} 

    .gr-button:hover {color: white !important;} 

    .gr-button {color: black !important;} 

    .gr-textbox {color: black !important;}

    .gr-Tab {color: black !important;}  /* Tab text color set to black */

""") as iface:
    
    # Title
    gr.Markdown("<h1 style='text-align: center; color:black;'>TransArt - Multimodal Application</h1>")

    # First Tab: Audio to Text -> Image
    with gr.Tab("Audio to Text"):
        gr.Markdown("<h3 style='text-align: center; color:black;'>Upload audio file, translate and generate an image</h3>")
        
        # Audio input and processing button
        with gr.Row():
            audio_input = gr.Audio(type="filepath", label="Upload Audio File")
            generate_image_checkbox = gr.Checkbox(label="Generate Image", value=False)
        
        # Outputs for transcription, translation, and image
        outputs = [
            gr.Textbox(label="Tamil Transcription"),
            gr.Textbox(label="English Translation"),
            gr.Image(label="Generated Image")  # Expecting an image output
        ]
        
        # Button for processing audio
        btn = gr.Button("Proceed Audio", elem_classes="btn-red")
        # Bind the correct function that returns transcription, translation, and an image
        btn.click(fn=process_audio, inputs=[audio_input, generate_image_checkbox], outputs=outputs)
    
    # Second Tab: Direct Prompt to Image Generation
    with gr.Tab("Prompt to Image"):
        gr.Markdown("<h3 style='text-align: center; color:black;'>Input a prompt and generate an image</h3>")
        
        # Text input for the prompt
        prompt_input = gr.Textbox(label="Enter Prompt", placeholder="Enter the scene description here...", lines=5)

        # Image output
        image_output = gr.Image(label="Generated Image")  # Expecting an image output

        # Button for generating the image from the prompt
        btn_image = gr.Button("Proceed Image Generation", elem_classes="btn-red")
        # Bind the correct function that returns an image
        btn_image.click(fn=generate_image_from_prompt, inputs=prompt_input, outputs=image_output)

# Launch the interface
iface.launch(server_name="0.0.0.0")