Update app.py

app.py

@@ -1,154 +1,139 @@
-import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+# Import libraries
 import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
+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"
 }
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
+
+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,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+            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.")