Update app.py

app.py

@@ -21,23 +21,60 @@ from safetensors.torch import load_file
 
 
 # Replace with your actual API key
-api_key = "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"
-client = Groq(api_key=api_key)
+os.environ['hugging']
+H_key = os.getenv('hugging')
+API_URL = "https://api-inference.huggingface.co/models/Artples/LAI-ImageGeneration-vSDXL-2"
+headers = {"Authorization": f"Bearer {H_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
+# 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")
+
+#key groq
+os.environ['groq']
+api_key = os.getenv('groq')
+client = Groq(api_key=api_key)
+
+def query(payload, max_retries=5):
+    for attempt in range(max_retries):
+        response = requests.post(API_URL, headers=headers, json=payload)
 
-# 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")
+        if response.status_code == 503:
+            print(f"Model is still loading, retrying... Attempt {attempt + 1}/{max_retries}")
+            estimated_time = min(response.json().get("estimated_time", 60), 60)
+            time.sleep(estimated_time)
+            continue
 
-os.environ['HF_API_KEY']
-api_key = os.getenv('HF_API_KEY')
-API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
+        if response.status_code != 200:
+            print(f"Error: Received status code {response.status_code}")
+            print(f"Response: {response.text}")
+            return None
+
+        return response.content
+
+    print(f"Failed to generate image after {max_retries} attempts.")
+    return None
+
+def generate_image_from_prompt(prompt):
+    image_bytes = query({"inputs": prompt})
+
+    if image_bytes is None:
+        return None
+
+    try:
+        image = Image.open(io.BytesIO(image_bytes))  # Opening the image from bytes
+        return image
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
 
 # Function to transcribe, translate, and generate an image
 def process_audio(audio_path, generate_image):
@@ -68,7 +105,8 @@ def process_audio(audio_path, generate_image):
     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]
+            #img = pipe(translation, num_inference_steps=4, guidance_scale=0).images[0]
+            img=generate_image_from_prompt(translation)
             return tamil_text, translation, img
         except Exception as e:
             return tamil_text, translation, f"An error occurred during image generation: {str(e)}"
@@ -76,48 +114,27 @@ def process_audio(audio_path, generate_image):
     return tamil_text, translation, None
 
 
-def query(payload, max_retries=5):
-    for attempt in range(max_retries):
-        response = requests.post(API_URL, headers=headers, json=payload)
-
-        if response.status_code == 503:
-            print(f"Model is still loading, retrying... Attempt {attempt + 1}/{max_retries}")
-            estimated_time = min(response.json().get("estimated_time", 60), 60)
-            time.sleep(estimated_time)
-            continue
+def chatbox(prompt):
+    try:
+        chat_completion = client.chat.completions.create(
+                messages=[{"role": "user", "content": prompt}],
+                model="llama-3.2-90b-text-preview"
+            )
+        chatbot_response = chat_completion.choices[0].message.content
 
-        if response.status_code != 200:
-            print(f"Error: Received status code {response.status_code}")
-            print(f"Response: {response.text}")
-            return None
+    except Exception as e:
+        return f"An error occurred during chatbot interaction: {str(e)}", None
 
-        return response.content
+    try:   
+        img=generate_image_from_prompt(prompt)
+    except Exception as e:
+            return chatbot_response, None
 
-    print(f"Failed to generate image after {max_retries} attempts.")
-    return None
+    return chatbot_response, img
     
 # Function for direct prompt to image generation
-def generate_image_from_prompt(prompt):
-    image_bytes = query({"inputs": prompt})
 
-    if image_bytes is None:
-        error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
-        d = ImageDraw.Draw(error_img)
-        d.text((10, 150), "Image Generation Failed", fill=(255, 255, 255))
-        return error_img
 
-    # Debug: Check the type of image_bytes
-    print(f"Received image_bytes type: {type(image_bytes)}")
-
-    try:
-        image = Image.open(io.BytesIO(image_bytes))
-        return image
-    except Exception as e:
-        print(f"Error while opening image: {e}")
-        error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
-        d = ImageDraw.Draw(error_img)
-        d.text((10, 150), "Invalid Image Data", fill=(255, 255, 255))
-        return error_img
 
 
 # Assuming your 'process_audio' and 'generate_image_from_prompt' functions are defined elsewhere
@@ -171,5 +188,21 @@ with gr.Blocks(css="""
         # Bind the correct function that returns an image
         btn_image.click(fn=generate_image_from_prompt, inputs=prompt_input, outputs=image_output)
 
+    #third tab: Direct prompt
+    with gr.Tab("Chatbot - image generation"):
+        gr.Markdown("<h2 style='text-align: center; color:black;'>Input a prompt and generate an image</h2>")
+        
+        prompt_input=gr.Textbox(label="Enter Prompt", placeholder="Enter the scene description here...", lines=2)
+         # Image output
+        output = [
+            gr.Textbox(label="Chatbot - response"),
+            gr.Image(label="Generated Image")  # Expecting an image output
+        ]
+          # Expecting an image output
+        # chatbox_output = 
+        btn_image = gr.Button("Chatbot Response Generation", elem_classes="btn-red")
+        # Bind the correct function that returns an image
+        btn_image.click(fn=chatbox, inputs=prompt_input, outputs=output)
+        
 # Launch the interface
 iface.launch(server_name="0.0.0.0")