import gradio as gr
import requests
import os
from PIL import Image
from io import BytesIO
from tqdm import tqdm
import time
import random

# Defining the repository information and the trigger word
repo = "stabilityai/stable-diffusion-xl-base-1.0"
trigger_word = "T shirt design, TshirtDesignAF, "

# Directory to save images
output_dir = "saved_designs"
os.makedirs(output_dir, exist_ok=True)  # Create directory if it doesn't exist

def generate_images(prompt):
    print("Generating 10 unique images with prompt:", prompt)
    api_url = f"https://api-inference.huggingface.co/models/{repo}"
    #token = os.getenv("API_TOKEN")  # Uncomment and use your Hugging Face API token
    headers = {
        #"Authorization": f"Bearer {token}"
    }
    
    images = []
    for i in range(10):
        # Add a unique seed to each prompt to ensure different images
        unique_seed = random.randint(1000, 9999)
        full_prompt = f"{prompt} {trigger_word} seed_{unique_seed}"

        payload = {
            "inputs": full_prompt,
            "parameters": {
                "negative_prompt": "(worst quality, low quality, normal quality, lowres, low details, oversaturated, undersaturated, overexposed, underexposed, grayscale, bw, bad photo, bad photography, bad art:1.4), (watermark, signature, text font, username, error, logo, words, letters, digits, autograph, trademark, name:1.2), (blur, blurry, grainy), morbid, ugly, asymmetrical, mutated malformed, mutilated, poorly lit, bad shadow, draft, cropped, out of frame, cut off, censored, jpeg artifacts, out of focus, glitch, duplicate, (airbrushed, cartoon, anime, semi-realistic, cgi, render, blender, digital art, manga, amateur:1.3), (3D ,3D Game, 3D Game Scene, 3D Character:1.1), (bad hands, bad anatomy, bad body, bad face, bad teeth, bad arms, bad legs, deformities:1.3)",
                "num_inference_steps": 30,
                "scheduler": "DPMSolverMultistepScheduler"
            },
        }

        error_count = 0
        pbar = tqdm(total=None, desc=f"Loading model {i+1}/10")
        while True:
            print(f"Sending request to API for image {i+1}...")
            response = requests.post(api_url, headers=headers, json=payload)
            print("API response status code:", response.status_code)
            if response.status_code == 200:
                print(f"Image {i+1} generation successful!")
                img = Image.open(BytesIO(response.content))
                images.append(img)
                
                # Save the image to the output directory
                img_filename = os.path.join(output_dir, f"{prompt.replace(' ', '_')}_design_{i+1}.png")
                img.save(img_filename)
                print(f"Image saved as: {img_filename}")
                break
            elif response.status_code == 503:
                time.sleep(1)
                pbar.update(1)
            elif response.status_code == 500 and error_count < 5:
                time.sleep(1)
                error_count += 1
            else:
                print(f"API Error for image {i+1}: {response.status_code}")
                raise Exception(f"API Error: {response.status_code}")

    return images

iface = gr.Interface(
    fn=generate_images,
    inputs=gr.Textbox(lines=2, placeholder="Type your prompt here..."),
    outputs=gr.Gallery(label="Generated Images", columns=5),  # Display images in a grid of 5 columns
    title="Design by rahul7star",
    description="Make designs for your clothes",
    examples=[["Cute Panda"], ["Skull"]]
)

print("Launching Gradio interface...")
iface.launch()