Upload app.py

app.py

@@ -0,0 +1,116 @@
+
+import gradio as gr
+import torch
+from diffusers import AutoPipelineForText2Image, DDIMScheduler
+from transformers import CLIPVisionModelWithProjection
+from diffusers.utils import load_image
+from PIL import Image
+import os
+import json
+import gc
+import traceback
+
+STYLE_MAP = {
+    "pixar": [
+        "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/style_ziggy/img0.png",
+        "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/style_ziggy/img1.png",
+        "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/style_ziggy/img2.png",
+        "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/style_ziggy/img3.png",
+        "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/style_ziggy/img4.png"
+    ]
+}
+
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"🚀 Device: {device}, torch_dtype: {torch_dtype}")
+
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+    "h94/IP-Adapter",
+    subfolder="models/image_encoder",
+    torch_dtype=torch_dtype,
+)
+
+pipeline = AutoPipelineForText2Image.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch_dtype,
+    image_encoder=image_encoder,
+    variant="fp16" if torch.cuda.is_available() else None
+).to(device)
+
+pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
+pipeline.load_ip_adapter(
+    "h94/IP-Adapter",
+    subfolder="sdxl_models",
+    weight_name=[
+        "ip-adapter-plus_sdxl_vit-h.safetensors",
+        "ip-adapter-plus-face_sdxl_vit-h.safetensors"
+    ]
+)
+pipeline.set_ip_adapter_scale([0.7, 0.3])
+pipeline.enable_model_cpu_offload()
+pipeline.enable_vae_tiling()
+
+def generate_storybook(data):
+    print("📥 Input JSON received:")
+    print(json.dumps(data, indent=2))
+
+    character_image_url = data["character_image_url"]
+    style = data["style"]
+    scenes = data["scenes"]
+
+    face_image = load_image(character_image_url)
+    style_images = [load_image(url) for url in STYLE_MAP.get(style, [])]
+
+    images = []
+
+    for i, prompt in enumerate(scenes):
+        print(f"🎬 Generating scene {i+1}: {prompt}")
+        try:
+            torch.cuda.empty_cache()
+            gc.collect()
+
+            result = pipeline(
+                prompt=prompt,
+                ip_adapter_image=[style_images, face_image],
+                negative_prompt="blurry, bad anatomy, low quality",
+                width=448,
+                height=672,
+                guidance_scale=5.0,
+                num_inference_steps=15,
+                generator=torch.Generator(device).manual_seed(i + 42)
+            )
+
+            image = result.images[0] if hasattr(result, "images") else result
+            print(f"🖼️ Image type: {type(image)}")
+
+            if isinstance(image, Image.Image):
+                images.append(image)
+                print(f"✅ Scene {i+1} added to image list.")
+            else:
+                print(f"⚠️ Scene {i+1} is not a valid image object.")
+
+        except Exception as e:
+            print(f"❌ Exception during scene {i+1}: {e}")
+            traceback.print_exc()
+
+    print(f"📦 Returning {len(images)} image(s)")
+    return images
+
+def generate_storybook_from_textbox(json_input_text):
+    try:
+        data = json.loads(json_input_text)
+        return generate_storybook(data)
+    except Exception as e:
+        print(f"❌ JSON parse or generation error: {e}")
+        traceback.print_exc()
+        return [f"Error: {str(e)}"]
+
+iface = gr.Interface(
+    fn=generate_storybook_from_textbox,
+    inputs=gr.Textbox(label="Input JSON", lines=20, placeholder="{...}"),
+    outputs=gr.Gallery(label="Generated Story Scenes", show_label=True, columns=1),
+    title="AI Storybook Generator (Low VRAM Mode)",
+    description="Optimized for lower VRAM GPUs. Paste JSON to generate consistent scenes."
+)
+
+iface.launch()