Spaces:
Sleeping
Sleeping
| # app.py - Improved Custom Content Generator | |
| import gradio as gr | |
| from transformers import pipeline, set_seed, AutoTokenizer, AutoModelForCausalLM | |
| from diffusers import StableDiffusionPipeline | |
| import torch | |
| # ---------- CONFIG ---------- | |
| TEXT_MODEL = "openai-community/gpt2" # Text model | |
| CODE_MODEL = "Salesforce/codegen-350M-multi" # Better for code | |
| IMAGE_MODEL = "stabilityai/stable-diffusion-2-1" # Higher quality images | |
| SEED = 42 | |
| # ---------------------------- | |
| # Load text generator | |
| text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL) | |
| text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL) | |
| text_pipe = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer, | |
| device_map="auto" if torch.cuda.is_available() else None) | |
| # Load code generator | |
| code_tokenizer = AutoTokenizer.from_pretrained(CODE_MODEL) | |
| code_model = AutoModelForCausalLM.from_pretrained(CODE_MODEL) | |
| code_pipe = pipeline("text-generation", model=code_model, tokenizer=code_tokenizer, | |
| device_map="auto" if torch.cuda.is_available() else None) | |
| # Load Stable Diffusion | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| try: | |
| sd_pipe = StableDiffusionPipeline.from_pretrained( | |
| IMAGE_MODEL, | |
| torch_dtype=torch.float16 if device == "cuda" else torch.float32 | |
| ).to(device) | |
| has_sd = True | |
| except Exception as e: | |
| sd_pipe = None | |
| has_sd = False | |
| print("Stable Diffusion not available:", e) | |
| set_seed(SEED) | |
| # ---------- Generation helpers ---------- | |
| def generate_text(prompt, max_len=150, temperature=0.7, top_k=50, top_p=0.95, num_return=1): | |
| out = text_pipe(prompt, max_length=max_len, do_sample=True, | |
| temperature=float(temperature), top_k=int(top_k), | |
| top_p=float(top_p), num_return_sequences=int(num_return)) | |
| return "\n\n===\n\n".join(o['generated_text'] for o in out) | |
| def generate_image(prompt, steps=35, guidance_scale=8.0, height=768, width=768): | |
| if not has_sd: | |
| return "Stable Diffusion model not loaded." | |
| return sd_pipe(prompt, num_inference_steps=int(steps), | |
| guidance_scale=float(guidance_scale), | |
| height=int(height), width=int(width)).images[0] | |
| def generate_code(prompt, language="Python", max_len=256): | |
| full_prompt = f"# {language} code\n# {prompt}\n" | |
| out = code_pipe(full_prompt, max_length=max_len, do_sample=True, | |
| temperature=0.2, top_k=50, top_p=0.95, num_return_sequences=1) | |
| return out[0]['generated_text'] | |
| # ---------- UI ---------- | |
| with gr.Blocks(title="Custom Content Generator") as demo: | |
| gr.Markdown("## Custom Content Generator — Text, Image, and Code.") | |
| with gr.Tabs(): | |
| with gr.TabItem("Text"): | |
| inp = gr.Textbox(label="Prompt", lines=3) | |
| max_len = gr.Slider(50, 512, value=150, step=10, label="Max length") | |
| temp = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature") | |
| top_k = gr.Slider(0, 200, value=50, step=1, label="Top-k") | |
| top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01, label="Top-p") | |
| out_text = gr.Textbox(label="Generated Text", lines=12) | |
| gr.Button("Generate Text").click( | |
| generate_text, [inp, max_len, temp, top_k, top_p, gr.Number(value=1, visible=False)], out_text | |
| ) | |
| with gr.TabItem("Image"): | |
| img_prompt = gr.Textbox(label="Image Prompt", lines=2) | |
| steps = gr.Slider(10, 50, value=35, step=1, label="Steps") | |
| guidance = gr.Slider(1.0, 20.0, value=8.0, step=0.1, label="Guidance scale") | |
| size = gr.Dropdown(choices=["512x512", "768x768"], value="768x768", label="Image size") | |
| img_out = gr.Image(label="Generated Image") | |
| def img_gen(prompt, steps, guidance, size): | |
| w, h = map(int, size.split("x")) | |
| return generate_image(prompt, steps=steps, guidance_scale=guidance, height=h, width=w) | |
| gr.Button("Generate Image").click(img_gen, [img_prompt, steps, guidance, size], img_out) | |
| with gr.TabItem("Code"): | |
| code_prompt = gr.Textbox(label="Describe Code Task", lines=3) | |
| lang = gr.Dropdown(choices=["Python", "JavaScript", "TypeScript", "Java", "C#", "Go"], value="Python") | |
| code_out = gr.Textbox(label="Generated Code", lines=18) | |
| gr.Button("Generate Code").click(generate_code, [code_prompt, lang], code_out) | |
| if __name__ == "__main__": | |
| demo.launch() | |