Update app.py

app.py

@@ -1,241 +1,88 @@
-import numpy as np
+# Шаг 1. Установка
+!pip install diffusers transformers accelerate safetensors gradio --quiet
+
+# Шаг 2. Импорт
 import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.distributions import Categorical
+from diffusers import StableDiffusionPipeline
+from PIL import Image, ImageFilter, ImageEnhance
+import numpy as np
 import gradio as gr
-import matplotlib.pyplot as plt
-import io
-from PIL import Image
 
-# ==========================================
-# 1. БЕСКОНЕЧНЫЙ МИР
-# ==========================================
-class InfiniteWorld:
-    CHUNK_SIZE = 16
-    VIEW_RADIUS = 8
-    
-    def __init__(self, seed=42):
-        self.seed = seed
-        self.chunks = {}
-        self.agent_pos = [0, 0]
-        self.steps = 0
-        self.max_steps = 1000
-        
-    def _get_chunk(self, cx, cy):
-        if (cx, cy) not in self.chunks:
-            rng = np.random.RandomState(hash((cx, cy, self.seed)) % (2**31))
-            chunk = np.zeros((self.CHUNK_SIZE, self.CHUNK_SIZE), dtype=np.float32)
-            noise = rng.rand(self.CHUNK_SIZE, self.CHUNK_SIZE)
-            chunk[noise > 0.7] = 1.0
-            self.chunks[(cx, cy)] = chunk
-        return self.chunks[(cx, cy)]
-    
-    def _world_coords(self, x, y):
-        cx, lx = divmod(x, self.CHUNK_SIZE)
-        cy, ly = divmod(y, self.CHUNK_SIZE)
-        return cx, cy, lx, ly
+# Шаг 3. Загрузка модели
+print("Загрузка модели...")
+model_id = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    safety_checker=None,
+    requires_safety_checker=False
+)
+pipe = pipe.to("cuda")
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
+print("Модель загружена!")
+
+# Шаг 4. Функция генерации
+def generate_image(prompt, negative_prompt):
+    # УСИЛЕНИЕ ПРОМПТА (если он короткий)
+    if len(prompt.split()) < 10:
+        prompt = prompt + ", detailed, sharp focus, masterpiece"
     
-    def get_block(self, x, y):
-        cx, cy, lx, ly = self._world_coords(x, y)
-        return self._get_chunk(cx, cy)[lx, ly]
+    generator = torch.Generator(device="cuda").manual_seed(2021)
     
-    def set_block(self, x, y, val):
-        cx, cy, lx, ly = self._world_coords(x, y)
-        self._get_chunk(cx, cy)[lx, ly] = val
+    with torch.autocast("cuda"):
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt if negative_prompt else "",
+            width=512,
+            height=768,
+            num_inference_steps=10,
+            guidance_scale=11.0,
+            generator=generator,
+            eta=0.8
+        ).images[0]
     
-    def reset(self):
-        self.agent_pos = [0, 0]
-        self.steps = 0
-        return self._get_obs()
+    # Постобработка
+    img_array = np.array(image)
+    noise = np.random.randint(-25, 25, img_array.shape, dtype=np.int16)
+    img_noisy = np.clip(img_array.astype(np.int16) + noise, 0, 255).astype(np.uint8)
     
-    def _get_obs(self):
-        x, y = self.agent_pos
-        patch = np.zeros((self.VIEW_RADIUS*2, self.VIEW_RADIUS*2, 3), dtype=np.float32)
-        for dx in range(-self.VIEW_RADIUS, self.VIEW_RADIUS):
-            for dy in range(-self.VIEW_RADIUS, self.VIEW_RADIUS):
-                wx, wy = x + dx, y + dy
-                block = self.get_block(wx, wy)
-                px = dx + self.VIEW_RADIUS
-                py = dy + self.VIEW_RADIUS
-                patch[px, py, 0] = block
-                patch[px, py, 1] = max(0, 1.0 - abs(dx)/self.VIEW_RADIUS)
-                patch[px, py, 2] = max(0, 1.0 - abs(dy)/self.VIEW_RADIUS)
-        patch[self.VIEW_RADIUS, self.VIEW_RADIUS, 1] = 1.0
-        return patch
+    img_final = Image.fromarray(img_noisy).filter(ImageFilter.GaussianBlur(radius=0.6))
     
-    def step(self, action):
-        self.steps += 1
-        reward = -0.005
-        done = self.steps >= self.max_steps
-        
-        if action == 0: self.agent_pos[0] -= 1
-        elif action == 1: self.agent_pos[0] += 1
-        elif action == 2: self.agent_pos[1] -= 1
-        elif action == 3: self.agent_pos[1] += 1
-        elif action == 4:
-            x, y = self.agent_pos
-            if self.get_block(x, y) == 0:
-                self.set_block(x, y, 1.0)
-                reward = 1.0
-        elif action == 5:
-            x, y = self.agent_pos
-            if self.get_block(x, y) == 1.0:
-                self.set_block(x, y, 0.0)
-                reward = 0.3
-        
-        return self._get_obs(), reward, done, {}
-
-
-# ==========================================
-# 2. PPO AGENT
-# ==========================================
-class PPOAgent(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.encoder = nn.Sequential(
-            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(),
-            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
-            nn.Conv2d(64, 64, 3, stride=2, padding=1), nn.ReLU(),
-            nn.AdaptiveAvgPool2d((4, 4)),
-            nn.Flatten()
-        )
-        self.gru = nn.GRUCell(64 * 4 * 4, 256)
-        self.actor = nn.Linear(256, 6)
-        self.critic = nn.Linear(256, 1)
+    enhancer = ImageEnhance.Contrast(img_final)
+    img_final = enhancer.enhance(1.3)
     
-    def forward(self, obs, hidden=None):
-        features = self.encoder(obs.permute(0, 3, 1, 2))
-        h = self.gru(features, hidden)
-        return self.actor(h), self.critic(h), h
+    enhancer = ImageEnhance.Color(img_final)
+    img_final = enhancer.enhance(1.3)
     
-    def act(self, obs, hidden=None):
-        with torch.no_grad():
-            logits, value, new_hidden = self.forward(obs.unsqueeze(0), hidden)
-            dist = Categorical(logits=logits)
-            action = dist.sample()
-            return action.item(), dist.log_prob(action), value.squeeze(), new_hidden
+    return img_final
 
-
-# ==========================================
-# 3. ОБУЧЕНИЕ
-# ==========================================
-def train_ppo(episodes=100):
-    env = InfiniteWorld()
-    agent = PPOAgent()
-    optimizer = optim.Adam(agent.parameters(), lr=3e-4)
+# Шаг 5. Интерфейс
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("#")
     
-    for ep in range(episodes):
-        obs = env.reset()
-        hidden = None
-        buffers = {'obs': [], 'actions': [], 'log_probs': [], 'rewards': [], 'values': []}
-        
-        for _ in range(256):
-            obs_t = torch.FloatTensor(obs)
-            action, log_prob, value, hidden = agent.act(obs_t, hidden)
-            next_obs, reward, done, _ = env.step(action)
-            
-            buffers['obs'].append(obs_t)
-            buffers['actions'].append(action)
-            buffers['log_probs'].append(log_prob)
-            buffers['rewards'].append(reward)
-            buffers['values'].append(value)
-            
-            obs = next_obs
-            if done:
-                obs = env.reset()
-                hidden = None
-        
-        # GAE
-        returns, advantages = [], []
-        R, A = 0, 0
-        for i in reversed(range(len(buffers['rewards']))):
-            R = buffers['rewards'][i] + 0.99 * R
-            next_val = buffers['values'][i+1].item() if i < len(buffers['values'])-1 else 0
-            delta = buffers['rewards'][i] + 0.99 * next_val - buffers['values'][i].item()
-            A = delta + 0.99 * 0.95 * A
-            returns.insert(0, R)
-            advantages.insert(0, A)
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt_input = gr.Textbox(
+                label="",
+                placeholder="Ваш промпт...",
+                lines=10
+            )
+            negative_input = gr.Textbox(
+                label="",
+                placeholder="Ваш негативный промпт...",
+                lines=5
+            )
+            generate_btn = gr.Button("Создать", variant="primary")
         
-        returns = torch.FloatTensor(returns)
-        advantages = torch.FloatTensor(advantages)
-        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
-        
-        obs_batch = torch.stack(buffers['obs'])
-        actions_batch = torch.LongTensor(buffers['actions'])
-        old_log_probs = torch.stack(buffers['log_probs']).detach()
-        
-        for _ in range(4):
-            logits, values, _ = agent.forward(obs_batch)
-            dist = Categorical(logits=logits)
-            new_log_probs = dist.log_prob(actions_batch)
-            ratio = (new_log_probs - old_log_probs).exp()
-            surr = torch.min(ratio * advantages, torch.clamp(ratio, 0.8, 1.2) * advantages)
-            loss = -surr.mean() + 0.5 * (returns - values.squeeze()).pow(2).mean() - 0.01 * dist.entropy().mean()
-            
-            optimizer.zero_grad()
-            loss.backward()
-            nn.utils.clip_grad_norm_(agent.parameters(), 0.5)
-            optimizer.step()
-        
-        if ep % 20 == 0:
-            print(f"Ep {ep} | Chunks: {len(env.chunks)}")
+        with gr.Column(scale=2):
+            output = gr.Image(label="", type="pil", height=600)
     
-    return agent
-
-
-# ==========================================
-# 4. ГРАФИЧЕСКИЙ ИНТЕРФЕЙС (ИСПРАВЛЕНО)
-# ==========================================
-def fig_to_pil(fig):
-    """Конвертирует matplotlib figure в PIL Image без schema-багов"""
-    buf = io.BytesIO()
-    fig.savefig(buf, format='png', bbox_inches='tight')
-    buf.seek(0)
-    img = Image.open(buf)
-    plt.close(fig)
-    return img
-
-
-def run_simulation(n_steps):
-    n_steps = int(n_steps)
-    agent = run_simulation.agent
-    env = InfiniteWorld(seed=np.random.randint(0, 99999))
-    obs = env.reset()
-    hidden = None
-    
-    images = []
-    with torch.no_grad():
-        for _ in range(min(n_steps, 300)):
-            fig, ax = plt.subplots(figsize=(4, 4))
-            ax.imshow(obs)
-            ax.set_title(f"Pos: {env.agent_pos}")
-            ax.axis('off')
-            images.append(fig_to_pil(fig))
-            
-            obs_t = torch.FloatTensor(obs)
-            action, _, _, hidden = agent.act(obs_t, hidden)
-            obs, _, done, _ = env.step(action)
-            if done:
-                break
-    
-    return images
-
-
-# Предобучаем модель один раз при загрузке
-print("🏗️ Обучение агента...")
-run_simulation.agent = train_ppo(episodes=80)
-run_simulation.agent.eval()
-print("✅ Обучение завершено!")
-
-
-# Интерфейс БЕЗ типизации возврата, БЕЗ Gallery
-with gr.Blocks(title="Infinite Builder") as demo:
-    gr.Markdown("# 🌍 Бесконечный мир: PPO-агент")
-    slider = gr.Slider(50, 300, value=100, step=50, label="Шагов")
-    btn = gr.Button("▶️ Запустить")
-    output = gr.Gallery(label="Результат", columns=4)
-    btn.click(fn=run_simulation, inputs=[slider], outputs=[output])
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt_input, negative_input],
+        outputs=output
+    )
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+# Шаг 6. Запуск
+demo.launch(share=True, debug=False)