File size: 5,079 Bytes
69a8ec7 8dbcb99 69a8ec7 8dbcb99 69a8ec7 8dbcb99 69a8ec7 |
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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
# Import libraries
import torch
import numpy as np
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from diffusers import StableDiffusionPipeline
from IPython.display import display
### --- STEP 1: Load TinyLlama for Text Generation --- ###
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
# Initialize text generation pipeline
comic_pipeline = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer
)
### --- STEP 2: Load Stable Diffusion XL for High-Quality Images --- ###
model_id = "stabilityai/sd-turbo" # Best for artistic comic style
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
pipe.to("cuda") # Move to GPU for better performance
### --- STEP 3: User Inputs a Prompt & Number of Panels --- ###
user_prompt = input("Enter a topic for the comic strip: ") # Example: "Government of India"
# Get number of panels from the user
while True:
try:
num_panels = int(input("Enter the number of comic panels (3 to 6): "))
if 3 <= num_panels <= 6:
break
else:
print("❌ Please enter a number between 3 and 6.")
except ValueError:
print("❌ Invalid input! Please enter a number between 3 and 6.")
### --- STEP 4: User Chooses an Art Style --- ###
art_styles = {
"1": "Classic Comic",
"2": "Anime",
"3": "Cartoon",
"4": "Noir",
"5": "Cyberpunk",
"6": "Watercolor"
}
print("\n🎨 Choose an Art Style for the Comic:")
for key, style in art_styles.items():
print(f"{key}. {style}")
while True:
art_choice = input("\nEnter the number for your preferred art style: ")
if art_choice in art_styles:
chosen_style = art_styles[art_choice]
print(f"✅ You selected: {chosen_style}")
break
else:
print("❌ Invalid choice! Please enter a valid number.")
### --- STEP 5: Generate Comic-Style Breakdown Using TinyLlama --- ###
instruction = (
f"Generate a structured {num_panels}-panel comic strip description for the topic. "
"Each panel should have a simple but clear scene description. "
"Keep it short and focus on visuals for easy image generation.\n\n"
"Topic: " + user_prompt + "\n\n"
"Comic Strip Panels:\n"
)
response = comic_pipeline(
instruction,
max_new_tokens=400, # Ensure full response
temperature=0.7,
repetition_penalty=1.1,
do_sample=True
)[0]['generated_text']
# Extract only the structured comic description
comic_breakdown = response.replace(instruction, "").strip()
comic_panels = [line.strip() for line in comic_breakdown.split("\n") if line.strip()][:num_panels]
print("\n🔹 Comic Strip Breakdown:\n", "\n".join(comic_panels)) # Show generated panels
### --- STEP 6: Generate High-Quality Comic-Style Images --- ###
def generate_comic_image(description, style):
"""
Generates a comic panel image using Stable Diffusion Turbo.
"""
# Validate style input (fallback to "Comic" if invalid)
valid_styles = ["Comic", "Anime", "Cyberpunk", "Watercolor", "Pixel Art"]
chosen_style = style if style in valid_styles else "Comic"
# Refined prompt (shorter, SD-Turbo-friendly)
prompt = f"{description}, {chosen_style} style, bold outlines, vibrant colors, dynamic action."
# Negative prompt (avoiding unwanted elements)
negative_prompt = "blurry, distorted, text, watermark, low quality, extra limbs, messy background"
try:
# Generate image with optimized parameters
image = pipe(
prompt,
negative_prompt=negative_prompt,
num_inference_steps=30, # Faster processing for SD-Turbo
guidance_scale=7
).images[0]
return image
except Exception as e:
print(f"❌ Error generating image: {e}")
return None # Return None if generation fails
# Generate images for each panel
comic_images = [generate_comic_image(panel, chosen_style) for panel in comic_panels]
# Remove None values if any images failed to generate
comic_images = [img for img in comic_images if img is not None]
if comic_images:
### --- STEP 7: Arrange Images in a Grid Based on Panel Count --- ###
grid_map = {3: (1, 3), 4: (2, 2), 5: (2, 3), 6: (2, 3)}
rows, cols = grid_map.get(len(comic_images), (1, len(comic_images)))
panel_width, panel_height = comic_images[0].size
comic_strip = Image.new("RGB", (panel_width * cols, panel_height * rows))
# Paste images in grid format
for i, img in enumerate(comic_images):
x_offset = (i % cols) * panel_width
y_offset = (i // cols) * panel_height
comic_strip.paste(img, (x_offset, y_offset))
# Display and save the comic strip
display(comic_strip)
comic_strip.save("comic_strip.png")
print("\n✅ Comic strip saved as 'comic_strip.png'")
else:
print("\n❌ No images were generated.") |