Spaces:
Sleeping
Sleeping
File size: 4,686 Bytes
4a0d8b9 |
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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import os
import pandas as pd
import torch
import gc
import re
import random
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from diffusers import StableDiffusionPipeline
import gradio as gr
# Initialize the text generation pipeline with the pre-quantized 8-bit model
model_name = 'HuggingFaceTB/SmolLM-1.7B-Instruct'
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=-1) # Use CPU
# Load the Stable Diffusion model
model_id = "stabilityai/stable-diffusion-2-1-base" # Smaller model
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe = pipe.to("cpu") # Use CPU
# Create a directory to save the generated images
output_dir = 'generated_images'
os.makedirs(output_dir, exist_ok=True)
os.chmod(output_dir, 0o777)
# Function to generate a detailed visual description prompt
def generate_description_prompt(user_prompt, user_examples):
prompt = f'generate enclosed in quotes in the format "<description>" description according to guidelines of {user_prompt} different from {user_examples}'
try:
generated_text = text_generator(prompt, max_length=150, num_return_sequences=1, truncation=True)[0]['generated_text']
match = re.search(r'"(.*?)"', generated_text)
if match:
generated_description = match.group(1).strip() # Capture the description between quotes
return f'"{generated_description}"'
else:
return None
except Exception as e:
print(f"Error generating description for prompt '{user_prompt}': {e}")
return None
# Seed words pool
seed_words = []
used_words = set()
def generate_description(user_prompt, user_examples_list):
seed_words.extend(user_examples_list)
# Select a subject that has not been used
available_subjects = [word for word in seed_words if word not in used_words]
if not available_subjects:
print("No more available subjects to use.")
return None, None
subject = random.choice(available_subjects)
generated_description = generate_description_prompt(user_prompt, subject)
if generated_description:
# Remove any offending symbols
clean_description = generated_description.encode('ascii', 'ignore').decode('ascii')
# Print the generated description to the command line
print(f"Generated description for subject '{subject}': {clean_description}")
# Update used words and seed words
used_words.add(subject)
seed_words.append(clean_description.strip('"')) # Add the generated description to the seed bank array without quotes
return clean_description, subject
else:
return None, None
# Function to generate an image based on the description
def generate_image(description, seed=42):
prompt = f'detailed photorealistic full shot of {description}'
generator = torch.Generator().manual_seed(seed)
image = pipe(
prompt=prompt,
width=512,
height=512,
num_inference_steps=10, # Use 10 inference steps
generator=generator,
guidance_scale=7.5,
).images[0]
return image
# Gradio interface
def gradio_interface(user_prompt, user_examples):
user_examples_list = [example.strip().strip('"') for example in user_examples.split(',')]
generated_description, subject = generate_description(user_prompt, user_examples_list)
if generated_description:
# Generate image
image = generate_image(generated_description)
image_path = os.path.join(output_dir, f"image_{len(os.listdir(output_dir))}.png")
image.save(image_path)
os.chmod(image_path, 0o777)
return image, generated_description
else:
return None, "Failed to generate description."
iface = gr.Interface(
fn=gradio_interface,
inputs=[
gr.Textbox(lines=2, placeholder="Enter the generation task or general thing you are looking for"),
gr.Textbox(lines=2, placeholder='Provide a few examples (enclosed in quotes and separated by commas)')
],
outputs=[
gr.Image(label="Generated Image"),
gr.Textbox(label="Generated Description")
],
title="Description and Image Generator",
description="Generate detailed descriptions and images based on your input."
)
iface.launch(server_name="0.0.0.0", server_port=7860)
# Clear GPU memory when the process is closed
def clear_gpu_memory():
torch.cuda.empty_cache()
gc.collect()
print("GPU memory cleared.")
|