Spaces:
Sleeping
Sleeping
File size: 2,779 Bytes
73d6f9d d4a2ccc 73d6f9d |
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 |
import os
import gradio as gr
from groq import Groq
import torch
from diffusers import StableDiffusionPipeline
# ===== GROQ CLIENT =====
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
# ===== LOAD STABLE DIFFUSION =====
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16
)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
# ===== PROMPT ENHANCER USING GROQ =====
def enhance_prompt(user_prompt):
response = client.chat.completions.create(
model="llama-3.3-70b-versatile",
messages=[
{
"role": "system",
"content": (
"You are an expert AI image prompt engineer. "
"Enhance the user prompt with visual details, lighting, style, realism, "
"camera angle, and artistic quality. "
"Return ONLY the enhanced prompt."
)
},
{"role": "user", "content": user_prompt}
],
temperature=0.7,
max_tokens=150
)
return response.choices[0].message.content.strip()
# ===== IMAGE GENERATION =====
def generate_image(prompt):
if not prompt.strip():
return None, ""
enhanced_prompt = enhance_prompt(prompt)
image = pipe(
enhanced_prompt,
num_inference_steps=30,
guidance_scale=7.5
).images[0]
return image, enhanced_prompt
# ===== UI =====
with gr.Blocks(theme=gr.themes.Soft(), title="AI Image Generator (Groq + SD)") as demo:
gr.Markdown(
"""
# 🎨 AI Image Generator
**Groq LLM + Stable Diffusion**
Type a prompt → Get a stunning image
"""
)
prompt_input = gr.Textbox(
label="Image Prompt",
placeholder="e.g. A futuristic city at sunset",
lines=3
)
generate_btn = gr.Button("Generate Image 🚀")
image_output = gr.Image(label="Generated Image")
enhanced_prompt_output = gr.Textbox(
label="Enhanced Prompt (Generated by Groq)",
lines=4
)
generate_btn.click(
generate_image,
inputs=prompt_input,
outputs=[image_output, enhanced_prompt_output]
)
gr.Markdown("### 🔹 Demo Prompts (Click & Generate)")
demo_prompts = [
"A cyberpunk city at night",
"A realistic lion in the jungle",
"A futuristic robot chef",
"A fantasy castle in the clouds",
"A Pakistani truck art style painting"
]
for demo_text in demo_prompts:
gr.Button(demo_text).click(
generate_image,
inputs=gr.State(demo_text),
outputs=[image_output, enhanced_prompt_output]
)
demo.launch() |