Create tambah img2img - app.py

tambah img2img - app.py

@@ -0,0 +1,148 @@
+import torch
+import gradio as gr
+import gc
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+
+device = "cpu"
+
+current_pipe = None
+current_model = None
+history = []
+
+
+def load_model(model_choice):
+    global current_pipe, current_model
+
+    if current_model == model_choice and current_pipe is not None:
+        return current_pipe
+
+    if current_pipe is not None:
+        del current_pipe
+        gc.collect()
+
+    if model_choice == "SD 1.5":
+        model_id = "runwayml/stable-diffusion-v1-5"
+    else:
+        model_id = "Lykon/dreamshaper-8"
+
+    print(f"🔄 Loading model: {model_id}")
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
+    )
+
+    pipe.to(device)
+    pipe.enable_attention_slicing()
+    pipe.safety_checker = None
+
+    current_pipe = pipe
+    current_model = model_choice
+
+    return pipe
+
+
+# ✅ Upscale sederhana (2x)
+def upscale_image(image):
+    if image is None:
+        return None
+    w, h = image.size
+    return image.resize((w * 2, h * 2), Image.LANCZOS)
+
+
+# ✅ Generate function (txt2img + img2img)
+def generate(prompt, model_choice, steps, input_image, strength, do_upscale):
+    pipe = load_model(model_choice)
+
+    try:
+        full_prompt = prompt + ", masterpiece, ultra detailed, 4k"
+
+        if input_image is not None:
+            # img2img
+            image = pipe(
+                prompt=full_prompt,
+                image=input_image,
+                strength=strength,
+                num_inference_steps=steps,
+                guidance_scale=7
+            ).images[0]
+        else:
+            # txt2img
+            image = pipe(
+                prompt=full_prompt,
+                negative_prompt="blurry, low quality, bad anatomy",
+                num_inference_steps=steps,
+                guidance_scale=7,
+                height=512,
+                width=512
+            ).images[0]
+
+        # ✅ auto upscale
+        if do_upscale:
+            image = upscale_image(image)
+
+        # ✅ simpan ke history
+        history.append(image)
+
+        return image, history
+
+    except Exception as e:
+        return None, history
+
+
+# ✅ UI MODERN
+with gr.Blocks() as demo:
+
+    gr.Markdown("# 🎨 AI Image Generator (CPU Optimized)")
+    gr.Markdown("SD 1.5 & DreamShaper + Img2Img + Upscale + Gallery")
+
+    with gr.Row():
+        prompt = gr.Textbox(
+            label="Prompt",
+            placeholder="Contoh: a cyberpunk city at night",
+            lines=3
+        )
+
+    with gr.Row():
+        model_choice = gr.Radio(
+            choices=["SD 1.5", "DreamShaper"],
+            value="DreamShaper",
+            label="Model"
+        )
+
+        steps = gr.Slider(10, 30, value=20, step=1, label="Steps")
+
+    with gr.Row():
+        input_image = gr.Image(type="pil", label="Upload Image (optional - img2img)")
+        
+        strength = gr.Slider(
+            0.1, 1.0,
+            value=0.5,
+            step=0.1,
+            label="Strength (img2img saja)"
+        )
+
+    with gr.Row():
+        upscale = gr.Checkbox(label="Auto Upscale 2x", value=False)
+        generate_btn = gr.Button("🚀 Generate")
+
+    output = gr.Image(type="pil", label="Hasil")
+
+    gr.Markdown("## 🖼️ History")
+    gallery = gr.Gallery(label="Hasil Sebelumnya", columns=3)
+
+    def run(prompt, model_choice, steps, input_image, strength, upscale):
+        img, hist = generate(prompt, model_choice, steps, input_image, strength, upscale)
+        return img, hist
+
+    generate_btn.click(
+        fn=run,
+        inputs=[prompt, model_choice, steps, input_image, strength, upscale],
+        outputs=[output, gallery]
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()