import spaces
import torch
from diffusers import StableDiffusionPipeline, UNet2DConditionModel
import gradio as gr

model_id = "UX4567/Stable-Diffusion-v1-3.0"
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Model load ho raha hai...")

trained_unet = UNet2DConditionModel.from_pretrained(
    model_id,
    torch_dtype=torch.float16
)

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    unet=trained_unet,
    torch_dtype=torch.float16,
    safety_checker=None # thoda fast hoga
).to(device)

@spaces.GPU
def generate_image(prompt):
    if not prompt.strip():
        return None

    # Quality booster
    enhanced_prompt = f"{prompt}, ultra sharp, highly detailed, cinematic lighting, 8k"
    negative_prompt = "blurry, low quality, low-res, distorted, ugly, bad anatomy, watermark, text, deformed"

    image = pipe(
        prompt=enhanced_prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=35,
        guidance_scale=8.5, # 7.5 se 8.5 - thoda zyada sharp
        width=512,
        height=768 # tere cyberpunk wale ke liye vertical better lagega
    ).images[0]
    return image

# Better UI - tere screenshot jaisa
with gr.Blocks() as demo:
    gr.Markdown("# Calico Image Generator\nApne fine-tuned Stable Diffusion model se creative images generate karein!")

    with gr.Group():
        prompt_box = gr.Textbox(label="Enter your prompt", value="a futuristic cyberpunk city at night", lines=2)

    with gr.Row():
        clear_btn = gr.Button("Clear")
        submit_btn = gr.Button("Submit", variant="primary")

    output_img = gr.Image(label="Generated Image")

    gr.Markdown("Share via Link se share kar sakta hai tu!")

    submit_btn.click(fn=generate_image, inputs=prompt_box, outputs=output_img)
    clear_btn.click(lambda: ("", None), inputs=None, outputs=[prompt_box, output_img])

demo.launch()